xref: /petsc/src/dm/tutorials/ex22.c (revision 9371c9d470a9602b6d10a8bf50c9b2280a79e45a)
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   DM                 da3D;                            /* 3D DMDA object */
18   DM                 da2D;                            /* 2D DMDA object */
19   Vec                vec_full;                        /* Parent vector */
20   Vec                vec_extracted;                   /* Extracted slice vector (in DMDA ordering) */
21   Vec                vec_slice;                       /* vector in natural ordering */
22   Vec                vec_slice_g;                     /* aliased vector in natural ordering */
23   IS                 patchis_3d;                      /* IS to select slice and extract subvector */
24   IS                 patchis_2d;                      /* Patch IS for 2D vector, will be converted to application ordering */
25   IS                 scatis_extracted_slice;          /* PETSc indexed IS for extracted slice */
26   IS                 scatis_natural_slice;            /* natural/application ordered IS for slice*/
27   IS                 scatis_natural_slice_g;          /* aliased natural/application ordered IS  for slice */
28   VecScatter         vscat;                           /* scatter slice in DMDA ordering <-> slice in column major ordering */
29   AO                 da2D_ao;                         /* AO associated with 2D DMDA */
30   MPI_Comm           subset_mpi_comm = MPI_COMM_NULL; /* MPI communicator where the slice lives */
31   PetscScalar     ***vecdata3d;                       /* Pointer to access 3d parent vector */
32   const PetscScalar *array;                           /* pointer to create aliased Vec */
33   PetscInt           Mx = 4, My = 4, Mz = 4;          /* Dimensions for 3D DMDA */
34   const PetscInt    *l1, *l2;                         /* 3D DMDA layout */
35   PetscInt           M1 = -1, M2 = -1;                /* Dimensions for 2D DMDA */
36   PetscInt           m1 = -1, m2 = -1;                /* Layouts for 2D DMDA */
37   PetscInt           gp        = 2;                   /* grid point along sliceaxis to pick the slice */
38   DMDirection        sliceaxis = DM_X;                /* Select axis along which the slice will be extracted */
39   PetscInt           i, j, k;                         /* Iteration indices */
40   PetscInt           ixs, iys, izs;                   /* Corner indices for 3D vector */
41   PetscInt           ixm, iym, izm;                   /* Widths of parent vector */
42   PetscInt           low, high;                       /* ownership range indices */
43   PetscInt           in;                              /* local size index for IS*/
44   PetscInt           vn;                              /* local size index */
45   const PetscInt    *is_array;                        /* pointer to create aliased IS */
46   MatStencil         lower, upper;                    /* Stencils to select slice for Vec */
47   PetscBool          patchis_offproc = PETSC_FALSE;   /* flag to DMDACreatePatchIS indicating that off-proc values are to be ignored */
48   PetscMPIInt        rank, size;                      /* MPI rank and size */
49 
50   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
51      Initialize program and set problem parameters
52      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
53   PetscFunctionBeginUser;
54   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
55   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
56   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
57 
58   PetscOptionsBegin(PETSC_COMM_WORLD, "", "ex22 DMDA tutorial example options", "DMDA");
59   PetscCall(PetscOptionsRangeInt("-Mx", "dimension along x-axis", "ex22.c", Mx, &Mx, NULL, 0, PETSC_MAX_INT));
60   PetscCall(PetscOptionsRangeInt("-My", "dimension along y-axis", "ex22.c", My, &My, NULL, 0, PETSC_MAX_INT));
61   PetscCall(PetscOptionsRangeInt("-Mz", "dimension along z-axis", "ex22.c", Mz, &Mz, NULL, 0, PETSC_MAX_INT));
62   PetscCall(PetscOptionsEnum("-sliceaxis", "axis along which 2D slice is extracted from : X, Y, Z", "", sliceaxes, (PetscEnum)sliceaxis, (PetscEnum *)&sliceaxis, NULL));
63   PetscCall(PetscOptionsRangeInt("-gp", "index along sliceaxis at which 2D slice is extracted", "ex22.c", gp, &gp, NULL, 0, PETSC_MAX_INT));
64   PetscOptionsEnd();
65 
66   /* Ensure that the requested slice is not out of bounds for the selected axis */
67   if (sliceaxis == DM_X) {
68     PetscCheck(gp <= Mx, PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
69   } else if (sliceaxis == DM_Y) {
70     PetscCheck(gp <= My, PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
71   } else if (sliceaxis == DM_Z) {
72     PetscCheck(gp <= Mz, PETSC_COMM_WORLD, PETSC_ERR_USER_INPUT, "grid point along sliceaxis is larger than largest index!");
73   }
74 
75   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
76      Create 3D DMDA object.
77      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
78   PetscCall(DMDACreate3d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, Mx, My, Mz, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, 1, 1, NULL, NULL, NULL, &da3D));
79   PetscCall(DMSetFromOptions(da3D));
80   PetscCall(DMSetUp(da3D));
81 
82   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
83      Create the parent vector
84      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
85   PetscCall(DMCreateGlobalVector(da3D, &vec_full));
86   PetscCall(PetscObjectSetName((PetscObject)vec_full, "full_vector"));
87 
88   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
89      Populate the 3D vector
90      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
91   PetscCall(DMDAGetCorners(da3D, &ixs, &iys, &izs, &ixm, &iym, &izm));
92   PetscCall(DMDAVecGetArray(da3D, vec_full, &vecdata3d));
93   for (k = izs; k < izs + izm; k++) {
94     for (j = iys; j < iys + iym; j++) {
95       for (i = ixs; i < ixs + ixm; i++) { vecdata3d[k][j][i] = ((i - Mx / 2.0) * (j + Mx / 2.0)) + k * 100; }
96     }
97   }
98   PetscCall(DMDAVecRestoreArray(da3D, vec_full, &vecdata3d));
99 
100   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
101      Get an IS corresponding to a 2D slice
102      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
103   if (sliceaxis == DM_X) {
104     lower.i = gp;
105     lower.j = 0;
106     lower.k = 0;
107     upper.i = gp;
108     upper.j = My;
109     upper.k = Mz;
110   } else if (sliceaxis == DM_Y) {
111     lower.i = 0;
112     lower.j = gp;
113     lower.k = 0;
114     upper.i = Mx;
115     upper.j = gp;
116     upper.k = Mz;
117   } else if (sliceaxis == DM_Z) {
118     lower.i = 0;
119     lower.j = 0;
120     lower.k = gp;
121     upper.i = Mx;
122     upper.j = My;
123     upper.k = gp;
124   }
125   PetscCall(DMDACreatePatchIS(da3D, &lower, &upper, &patchis_3d, patchis_offproc));
126   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n IS to select slice from 3D DMDA vector : \n"));
127   PetscCall(ISView(patchis_3d, PETSC_VIEWER_STDOUT_WORLD));
128 
129   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
130      Use the obtained IS to extract the slice as a subvector
131      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
132   PetscCall(VecGetSubVector(vec_full, patchis_3d, &vec_extracted));
133 
134   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
135      View the extracted subvector
136      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
137   PetscCall(PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD, PETSC_VIEWER_ASCII_DENSE));
138   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in DMDA ordering : \n"));
139   PetscCall(VecView(vec_extracted, PETSC_VIEWER_STDOUT_WORLD));
140   PetscCall(PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD));
141 
142   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
143      Query 3D DMDA layout, get the subset MPI communicator
144      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
145   if (sliceaxis == DM_X) {
146     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL));
147     PetscCall(DMDAGetOwnershipRanges(da3D, NULL, &l1, &l2));
148     M1 = My;
149     M2 = Mz;
150     PetscCall(DMDAGetProcessorSubset(da3D, DM_X, gp, &subset_mpi_comm));
151   } else if (sliceaxis == DM_Y) {
152     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, NULL, &m2, NULL, NULL, NULL, NULL, NULL, NULL));
153     PetscCall(DMDAGetOwnershipRanges(da3D, &l1, NULL, &l2));
154     M1 = Mx;
155     M2 = Mz;
156     PetscCall(DMDAGetProcessorSubset(da3D, DM_Y, gp, &subset_mpi_comm));
157   } else if (sliceaxis == DM_Z) {
158     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL, NULL));
159     PetscCall(DMDAGetOwnershipRanges(da3D, &l1, &l2, NULL));
160     M1 = Mx;
161     M2 = My;
162     PetscCall(DMDAGetProcessorSubset(da3D, DM_Z, gp, &subset_mpi_comm));
163   }
164 
165   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
166      Create 2D DMDA object,
167      vector (that will hold the slice as a column major flattened array) &
168      index set (that will be used for scattering to the column major
169      indexed slice vector)
170      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
171   if (subset_mpi_comm != MPI_COMM_NULL) {
172     PetscCallMPI(MPI_Comm_size(subset_mpi_comm, &size));
173     PetscCall(PetscSynchronizedPrintf(subset_mpi_comm, "subset MPI subcomm size is : %d, includes global rank : %d \n", size, rank));
174     PetscCall(PetscSynchronizedFlush(subset_mpi_comm, PETSC_STDOUT));
175     PetscCall(DMDACreate2d(subset_mpi_comm, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, M1, M2, m1, m2, 1, 1, l1, l2, &da2D));
176     PetscCall(DMSetFromOptions(da2D));
177     PetscCall(DMSetUp(da2D));
178 
179     /* Create a 2D patch IS for the slice */
180     lower.i = 0;
181     lower.j = 0;
182     upper.i = M1;
183     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