xref: /petsc/src/dm/tutorials/ex22.c (revision d5b43468fb8780a8feea140ccd6fa3e6a50411cc)
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 
51   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
52      Initialize program and set problem parameters
53      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
54   PetscFunctionBeginUser;
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   PetscOptionsBegin(PETSC_COMM_WORLD, "", "ex22 DMDA tutorial example options", "DMDA");
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   PetscOptionsEnd();
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   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));
80   PetscCall(DMSetFromOptions(da3D));
81   PetscCall(DMSetUp(da3D));
82 
83   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
84      Create the parent vector
85      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
86   PetscCall(DMCreateGlobalVector(da3D, &vec_full));
87   PetscCall(PetscObjectSetName((PetscObject)vec_full, "full_vector"));
88 
89   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
90      Populate the 3D vector
91      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
92   PetscCall(DMDAGetCorners(da3D, &ixs, &iys, &izs, &ixm, &iym, &izm));
93   PetscCall(DMDAVecGetArray(da3D, vec_full, &vecdata3d));
94   for (k = izs; k < izs + izm; k++) {
95     for (j = iys; j < iys + iym; j++) {
96       for (i = ixs; i < ixs + ixm; i++) vecdata3d[k][j][i] = ((i - Mx / 2.0) * (j + Mx / 2.0)) + k * 100;
97     }
98   }
99   PetscCall(DMDAVecRestoreArray(da3D, vec_full, &vecdata3d));
100 
101   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
102      Get an IS corresponding to a 2D slice
103      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
104   if (sliceaxis == DM_X) {
105     lower.i = gp;
106     lower.j = 0;
107     lower.k = 0;
108     upper.i = gp;
109     upper.j = My;
110     upper.k = Mz;
111   } else if (sliceaxis == DM_Y) {
112     lower.i = 0;
113     lower.j = gp;
114     lower.k = 0;
115     upper.i = Mx;
116     upper.j = gp;
117     upper.k = Mz;
118   } else if (sliceaxis == DM_Z) {
119     lower.i = 0;
120     lower.j = 0;
121     lower.k = gp;
122     upper.i = Mx;
123     upper.j = My;
124     upper.k = gp;
125   }
126   PetscCall(DMDACreatePatchIS(da3D, &lower, &upper, &patchis_3d, patchis_offproc));
127   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n IS to select slice from 3D DMDA vector : \n"));
128   PetscCall(ISView(patchis_3d, PETSC_VIEWER_STDOUT_WORLD));
129 
130   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
131      Use the obtained IS to extract the slice as a subvector
132      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
133   PetscCall(VecGetSubVector(vec_full, patchis_3d, &vec_extracted));
134 
135   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
136      View the extracted subvector
137      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
138   PetscCall(PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD, PETSC_VIEWER_ASCII_DENSE));
139   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in DMDA ordering : \n"));
140   PetscCall(VecView(vec_extracted, PETSC_VIEWER_STDOUT_WORLD));
141   PetscCall(PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD));
142 
143   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
144      Query 3D DMDA layout, get the subset MPI communicator
145      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
146   if (sliceaxis == DM_X) {
147     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL));
148     PetscCall(DMDAGetOwnershipRanges(da3D, NULL, &l1, &l2));
149     M1 = My;
150     M2 = Mz;
151     PetscCall(DMDAGetProcessorSubset(da3D, DM_X, gp, &subset_mpi_comm));
152   } else if (sliceaxis == DM_Y) {
153     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, NULL, &m2, NULL, NULL, NULL, NULL, NULL, NULL));
154     PetscCall(DMDAGetOwnershipRanges(da3D, &l1, NULL, &l2));
155     M1 = Mx;
156     M2 = Mz;
157     PetscCall(DMDAGetProcessorSubset(da3D, DM_Y, gp, &subset_mpi_comm));
158   } else if (sliceaxis == DM_Z) {
159     PetscCall(DMDAGetInfo(da3D, NULL, NULL, NULL, NULL, &m1, &m2, NULL, NULL, NULL, NULL, NULL, NULL, NULL));
160     PetscCall(DMDAGetOwnershipRanges(da3D, &l1, &l2, NULL));
161     M1 = Mx;
162     M2 = My;
163     PetscCall(DMDAGetProcessorSubset(da3D, DM_Z, gp, &subset_mpi_comm));
164   }
165 
166   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
167      Create 2D DMDA object,
168      vector (that will hold the slice as a column major flattened array) &
169      index set (that will be used for scattering to the column major
170      indexed slice vector)
171      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
172   if (subset_mpi_comm != MPI_COMM_NULL) {
173     PetscCallMPI(MPI_Comm_size(subset_mpi_comm, &size));
174     PetscCall(PetscSynchronizedPrintf(subset_mpi_comm, "subset MPI subcomm size is : %d, includes global rank : %d \n", size, rank));
175     PetscCall(PetscSynchronizedFlush(subset_mpi_comm, PETSC_STDOUT));
176     PetscCall(DMDACreate2d(subset_mpi_comm, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, M1, M2, m1, m2, 1, 1, l1, l2, &da2D));
177     PetscCall(DMSetFromOptions(da2D));
178     PetscCall(DMSetUp(da2D));
179 
180     /* Create a 2D patch IS for the slice */
181     lower.i = 0;
182     lower.j = 0;
183     upper.i = M1;
184     upper.j = M2;
185     PetscCall(DMDACreatePatchIS(da2D, &lower, &upper, &patchis_2d, patchis_offproc));
186 
187     /* Convert the 2D patch IS to natural indexing (column major flattened) */
188     PetscCall(ISDuplicate(patchis_2d, &scatis_natural_slice));
189     PetscCall(DMDAGetAO(da2D, &da2D_ao));
190     PetscCall(AOPetscToApplicationIS(da2D_ao, scatis_natural_slice));
191     PetscCall(ISGetIndices(scatis_natural_slice, &is_array));
192     PetscCall(ISGetLocalSize(scatis_natural_slice, &in));
193 
194     /* Create an aliased IS on the 3D DMDA's communicator */
195     PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, in, is_array, PETSC_USE_POINTER, &scatis_natural_slice_g));
196     PetscCall(ISRestoreIndices(scatis_natural_slice, &is_array));
197 
198     /* Create a 2D DMDA global vector */
199     PetscCall(DMCreateGlobalVector(da2D, &vec_slice));
200     PetscCall(PetscObjectSetName((PetscObject)vec_slice, "slice_vector_natural"));
201     PetscCall(VecGetLocalSize(vec_slice, &vn));
202     PetscCall(VecGetArrayRead(vec_slice, &array));
203 
204     /* Create an aliased version of the above on the 3D DMDA's communicator */
205     PetscCall(VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, vn, M1 * M2, array, &vec_slice_g));
206     PetscCall(VecRestoreArrayRead(vec_slice, &array));
207   } else {
208     /* Ranks not part of the subset MPI communicator provide no entries, but the routines for creating
209        the IS and Vec on the 3D DMDA's communicator still need to called, since they are collective routines */
210     PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 0, NULL, PETSC_USE_POINTER, &scatis_natural_slice_g));
211     PetscCall(VecCreateMPIWithArray(PETSC_COMM_WORLD, 1, 0, M1 * M2, NULL, &vec_slice_g));
212   }
213   PetscCall(PetscBarrier(NULL));
214 
215   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
216      Create IS that maps from the extracted slice vector
217      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
218   PetscCall(VecGetOwnershipRange(vec_extracted, &low, &high));
219   PetscCall(ISCreateStride(PETSC_COMM_WORLD, high - low, low, 1, &scatis_extracted_slice));
220 
221   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
222      Scatter extracted subvector -> natural 2D slice vector
223      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
224   PetscCall(VecScatterCreate(vec_extracted, scatis_extracted_slice, vec_slice_g, scatis_natural_slice_g, &vscat));
225   PetscCall(VecScatterBegin(vscat, vec_extracted, vec_slice_g, INSERT_VALUES, SCATTER_FORWARD));
226   PetscCall(VecScatterEnd(vscat, vec_extracted, vec_slice_g, INSERT_VALUES, SCATTER_FORWARD));
227 
228   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
229      View the natural 2D slice vector
230      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
231   PetscCall(PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD, PETSC_VIEWER_ASCII_DENSE));
232   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n Extracted slice vector, in natural ordering : \n"));
233   PetscCall(VecView(vec_slice_g, PETSC_VIEWER_STDOUT_WORLD));
234   PetscCall(PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD));
235 
236   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
237      Restore subvector
238      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
239   PetscCall(VecRestoreSubVector(vec_full, patchis_3d, &vec_extracted));
240 
241   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
242      Destroy data structures and exit.
243      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
244   PetscCall(VecDestroy(&vec_full));
245   PetscCall(VecScatterDestroy(&vscat));
246   PetscCall(ISDestroy(&scatis_extracted_slice));
247   PetscCall(ISDestroy(&scatis_natural_slice_g));
248   PetscCall(VecDestroy(&vec_slice_g));
249   PetscCall(ISDestroy(&patchis_3d));
250   PetscCall(DMDestroy(&da3D));
251 
252   if (subset_mpi_comm != MPI_COMM_NULL) {
253     PetscCall(ISDestroy(&patchis_2d));
254     PetscCall(ISDestroy(&scatis_natural_slice));
255     PetscCall(VecDestroy(&vec_slice));
256     PetscCall(DMDestroy(&da2D));
257     PetscCallMPI(MPI_Comm_free(&subset_mpi_comm));
258   }
259 
260   PetscCall(PetscFinalize());
261   return 0;
262 }
263 
264 /*TEST
265 
266     test:
267       nsize: 1
268       args: -sliceaxis X -gp 0
269 
270     test:
271       suffix: 2
272       nsize:  2
273       args: -sliceaxis Y -gp 1
274       filter: grep -v "subset MPI subcomm"
275 
276     test:
277       suffix: 3
278       nsize:  3
279       args:  -sliceaxis Z -gp 2
280       filter: grep -v "subset MPI subcomm"
281 
282     test:
283       suffix: 4
284       nsize:  4
285       args: -sliceaxis X -gp 2
286       filter: grep -v "subset MPI subcomm"
287 
288     test:
289       suffix: 5
290       nsize:  4
291       args: -sliceaxis Z -gp 1
292       filter: grep -v "subset MPI subcomm"
293 
294 TEST*/
295