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