xref: /petsc/src/dm/tests/ex53.c (revision df4cd43f92eaa320656440c40edb1046daee8f75)
1 static char help[] = "Use DMDACreatePatchIS  to extract a slice from a vector, Command line options :\n\
2 mx/my/mz - set the dimensions of the parent vector\n\
3 dim - set the dimensionality of the parent vector (2,3)\n\
4 sliceaxis - Integer describing the axis along which the sice will be selected (0-X, 1-Y, 2-Z)\n\
5 sliceid - set the location where the slice will be extracted from the parent vector\n";
6 
7 /*
8    This test checks the functionality of DMDACreatePatchIS when
9    extracting a 2D vector from a 3D vector and 1D vector from a
10    2D vector.
11    */
12 
13 #include <petscdmda.h>
14 
15 int main(int argc, char **argv)
16 {
17   PetscMPIInt    rank, size;                    /* MPI rank and size */
18   PetscInt       mx = 4, my = 4, mz = 4;        /* Dimensions of parent vector */
19   PetscInt       sliceid   = 2;                 /* k (z) index to pick the slice */
20   PetscInt       sliceaxis = 2;                 /* Select axis along which the slice will be extracted */
21   PetscInt       dim       = 3;                 /* Dimension of the parent vector */
22   PetscInt       i, j, k;                       /* Iteration indices */
23   PetscInt       ixs, iys, izs;                 /* Corner indices for 3D vector */
24   PetscInt       ixm, iym, izm;                 /* Widths of parent vector */
25   PetscScalar ***vecdata3d;                     /* Pointer to access 3d parent vector */
26   PetscScalar  **vecdata2d;                     /* Pointer to access 2d parent vector */
27   DM             da;                            /* 2D/3D DMDA object */
28   Vec            vec_full;                      /* Parent vector */
29   Vec            vec_slice;                     /* Slice vector */
30   MatStencil     lower, upper;                  /* Stencils to select slice */
31   IS             selectis;                      /* IS to select slice and extract subvector */
32   PetscBool      patchis_offproc = PETSC_FALSE; /* flag to DMDACreatePatchIS indicating that off-proc values are to be ignored */
33 
34   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
35      Initialize program and set problem parameters
36      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
37   PetscFunctionBeginUser;
38   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
39   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
40   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
41 
42   PetscCall(PetscOptionsGetInt(NULL, NULL, "-mx", &mx, NULL));
43   PetscCall(PetscOptionsGetInt(NULL, NULL, "-my", &my, NULL));
44   PetscCall(PetscOptionsGetInt(NULL, NULL, "-mz", &mz, NULL));
45   PetscCall(PetscOptionsGetInt(NULL, NULL, "-dim", &dim, NULL));
46   PetscCall(PetscOptionsGetInt(NULL, NULL, "-sliceid", &sliceid, NULL));
47   PetscCall(PetscOptionsGetInt(NULL, NULL, "-sliceaxis", &sliceaxis, NULL));
48 
49   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
50      Create DMDA object.
51      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
52   if (dim == 3) {
53     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, &da));
54     PetscCall(DMSetFromOptions(da));
55     PetscCall(DMSetUp(da));
56   } else {
57     PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, mx, my, PETSC_DECIDE, PETSC_DECIDE, 1, 1, NULL, NULL, &da));
58     PetscCall(DMSetFromOptions(da));
59     PetscCall(DMSetUp(da));
60   }
61 
62   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
63      Create the parent vector
64      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*/
65   PetscCall(DMCreateGlobalVector(da, &vec_full));
66   PetscCall(PetscObjectSetName((PetscObject)vec_full, "full_vector"));
67 
68   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
69      Populate the 3D vector
70      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
71   PetscCall(DMDAGetCorners(da, &ixs, &iys, &izs, &ixm, &iym, &izm));
72   if (dim == 3) {
73     PetscCall(DMDAVecGetArray(da, vec_full, &vecdata3d));
74     for (k = izs; k < izs + izm; k++) {
75       for (j = iys; j < iys + iym; j++) {
76         for (i = ixs; i < ixs + ixm; i++) vecdata3d[k][j][i] = ((i - mx / 2) * (j + mx / 2)) + k * 100;
77       }
78     }
79     PetscCall(DMDAVecRestoreArray(da, vec_full, &vecdata3d));
80   } else {
81     PetscCall(DMDAVecGetArray(da, vec_full, &vecdata2d));
82     for (j = iys; j < iys + iym; j++) {
83       for (i = ixs; i < ixs + ixm; i++) vecdata2d[j][i] = ((i - mx / 2) * (j + mx / 2));
84     }
85     PetscCall(DMDAVecRestoreArray(da, vec_full, &vecdata2d));
86   }
87 
88   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
89      Get an IS corresponding to a 2D slice
90      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
91   if (sliceaxis == 0) {
92     lower.i = sliceid;
93     lower.j = 0;
94     lower.k = 0;
95     lower.c = 1;
96     upper.i = sliceid;
97     upper.j = my;
98     upper.k = mz;
99     upper.c = 1;
100   } else if (sliceaxis == 1) {
101     lower.i = 0;
102     lower.j = sliceid;
103     lower.k = 0;
104     lower.c = 1;
105     upper.i = mx;
106     upper.j = sliceid;
107     upper.k = mz;
108     upper.c = 1;
109   } else if (sliceaxis == 2 && dim == 3) {
110     lower.i = 0;
111     lower.j = 0;
112     lower.k = sliceid;
113     lower.c = 1;
114     upper.i = mx;
115     upper.j = my;
116     upper.k = sliceid;
117     upper.c = 1;
118   }
119   PetscCall(DMDACreatePatchIS(da, &lower, &upper, &selectis, patchis_offproc));
120   PetscCall(ISView(selectis, PETSC_VIEWER_STDOUT_WORLD));
121 
122   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
123      Use the obtained IS to extract the slice as a subvector
124      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
125   PetscCall(VecGetSubVector(vec_full, selectis, &vec_slice));
126 
127   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
128      View the extracted subvector
129      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
130   PetscCall(PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD, PETSC_VIEWER_ASCII_DENSE));
131   PetscCall(VecView(vec_slice, PETSC_VIEWER_STDOUT_WORLD));
132   PetscCall(PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD));
133 
134   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
135      Restore subvector, destroy data structures and exit.
136      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
137   PetscCall(VecRestoreSubVector(vec_full, selectis, &vec_slice));
138 
139   PetscCall(ISDestroy(&selectis));
140   PetscCall(DMDestroy(&da));
141   PetscCall(VecDestroy(&vec_full));
142 
143   PetscCall(PetscFinalize());
144   return 0;
145 }
146 
147 /*TEST
148 
149     test:
150       nsize: 1
151       args: -sliceaxis 0
152 
153     test:
154       suffix: 2
155       nsize:  2
156       args: -sliceaxis 1
157 
158     test:
159       suffix: 3
160       nsize:  4
161       args:  -sliceaxis 2
162 
163     test:
164       suffix: 4
165       nsize:  2
166       args: -sliceaxis 1 -dim 2
167 
168     test:
169       suffix: 5
170       nsize:  3
171       args: -sliceaxis 0 -dim 2
172 
173 TEST*/
174