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