xref: /petsc/src/dm/tutorials/swarm_ex3.c (revision 750b007cd8d816cecd9de99077bb0a703b4cf61a)
1 
2 static char help[] = "Tests DMSwarm with DMShell\n\n";
3 
4 #include <petscsf.h>
5 #include <petscdm.h>
6 #include <petscdmshell.h>
7 #include <petscdmda.h>
8 #include <petscdmswarm.h>
9 #include <petsc/private/dmimpl.h>
10 
11 PetscErrorCode _DMLocatePoints_DMDARegular_IS(DM dm, Vec pos, IS *iscell) {
12   PetscInt           p, n, bs, npoints, si, sj, milocal, mjlocal, mx, my;
13   DM                 dmregular;
14   PetscInt          *cellidx;
15   const PetscScalar *coor;
16   PetscReal          dx, dy;
17   PetscMPIInt        rank;
18 
19   PetscFunctionBegin;
20   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
21   PetscCall(VecGetLocalSize(pos, &n));
22   PetscCall(VecGetBlockSize(pos, &bs));
23   npoints = n / bs;
24 
25   PetscCall(PetscMalloc1(npoints, &cellidx));
26   PetscCall(DMGetApplicationContext(dm, &dmregular));
27   PetscCall(DMDAGetCorners(dmregular, &si, &sj, NULL, &milocal, &mjlocal, NULL));
28   PetscCall(DMDAGetInfo(dmregular, NULL, &mx, &my, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL));
29 
30   dx = 2.0 / ((PetscReal)mx);
31   dy = 2.0 / ((PetscReal)my);
32 
33   PetscCall(VecGetArrayRead(pos, &coor));
34   for (p = 0; p < npoints; p++) {
35     PetscReal coorx, coory;
36     PetscInt  mi, mj;
37 
38     coorx = PetscRealPart(coor[2 * p]);
39     coory = PetscRealPart(coor[2 * p + 1]);
40 
41     mi = (PetscInt)((coorx - (-1.0)) / dx);
42     mj = (PetscInt)((coory - (-1.0)) / dy);
43 
44     cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND;
45 
46     if ((mj >= sj) && (mj < sj + mjlocal)) {
47       if ((mi >= si) && (mi < si + milocal)) cellidx[p] = (mi - si) + (mj - sj) * milocal;
48     }
49     if (coorx < -1.0) cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND;
50     if (coorx > 1.0) cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND;
51     if (coory < -1.0) cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND;
52     if (coory > 1.0) cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND;
53   }
54   PetscCall(VecRestoreArrayRead(pos, &coor));
55   PetscCall(ISCreateGeneral(PETSC_COMM_SELF, npoints, cellidx, PETSC_OWN_POINTER, iscell));
56   PetscFunctionReturn(0);
57 }
58 
59 PetscErrorCode DMLocatePoints_DMDARegular(DM dm, Vec pos, DMPointLocationType ltype, PetscSF cellSF) {
60   IS              iscell;
61   PetscSFNode    *cells;
62   PetscInt        p, bs, npoints, nfound;
63   const PetscInt *boxCells;
64 
65   PetscFunctionBegin;
66   PetscCall(_DMLocatePoints_DMDARegular_IS(dm, pos, &iscell));
67   PetscCall(VecGetLocalSize(pos, &npoints));
68   PetscCall(VecGetBlockSize(pos, &bs));
69   npoints = npoints / bs;
70 
71   PetscCall(PetscMalloc1(npoints, &cells));
72   PetscCall(ISGetIndices(iscell, &boxCells));
73 
74   for (p = 0; p < npoints; p++) {
75     cells[p].rank  = 0;
76     cells[p].index = DMLOCATEPOINT_POINT_NOT_FOUND;
77     cells[p].index = boxCells[p];
78   }
79   PetscCall(ISRestoreIndices(iscell, &boxCells));
80   PetscCall(ISDestroy(&iscell));
81   nfound = npoints;
82   PetscCall(PetscSFSetGraph(cellSF, npoints, nfound, NULL, PETSC_OWN_POINTER, cells, PETSC_OWN_POINTER));
83   PetscCall(ISDestroy(&iscell));
84   PetscFunctionReturn(0);
85 }
86 
87 PetscErrorCode DMGetNeighbors_DMDARegular(DM dm, PetscInt *nneighbors, const PetscMPIInt **neighbors) {
88   DM dmregular;
89 
90   PetscFunctionBegin;
91   PetscCall(DMGetApplicationContext(dm, &dmregular));
92   PetscCall(DMGetNeighbors(dmregular, nneighbors, neighbors));
93   PetscFunctionReturn(0);
94 }
95 
96 PetscErrorCode SwarmViewGP(DM dms, const char prefix[]) {
97   PetscReal  *array;
98   PetscInt   *iarray;
99   PetscInt    npoints, p, bs;
100   FILE       *fp;
101   char        name[PETSC_MAX_PATH_LEN];
102   PetscMPIInt rank;
103 
104   PetscFunctionBegin;
105   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
106   PetscCall(PetscSNPrintf(name, PETSC_MAX_PATH_LEN - 1, "%s-rank%d.gp", prefix, rank));
107   fp = fopen(name, "w");
108   PetscCheck(fp, PETSC_COMM_SELF, PETSC_ERR_FILE_OPEN, "Cannot open file %s", name);
109   PetscCall(DMSwarmGetLocalSize(dms, &npoints));
110   PetscCall(DMSwarmGetField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array));
111   PetscCall(DMSwarmGetField(dms, "itag", NULL, NULL, (void **)&iarray));
112   for (p = 0; p < npoints; p++) fprintf(fp, "%+1.4e %+1.4e %1.4e\n", array[2 * p], array[2 * p + 1], (double)iarray[p]);
113   PetscCall(DMSwarmRestoreField(dms, "itag", NULL, NULL, (void **)&iarray));
114   PetscCall(DMSwarmRestoreField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array));
115   fclose(fp);
116   PetscFunctionReturn(0);
117 }
118 
119 /*
120  Create a DMShell and attach a regularly spaced DMDA for point location
121  Override methods for point location
122 */
123 PetscErrorCode ex3_1(void) {
124   DM          dms, dmcell, dmregular;
125   PetscMPIInt rank;
126   PetscInt    p, bs, nlocal, overlap, mx, tk;
127   PetscReal   dx;
128   PetscReal  *array, dt;
129   PetscInt   *iarray;
130   PetscRandom rand;
131 
132   PetscFunctionBegin;
133   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
134 
135   /* Create a regularly spaced DMDA */
136   mx      = 40;
137   overlap = 0;
138   PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_BOX, mx, mx, PETSC_DECIDE, PETSC_DECIDE, 1, overlap, NULL, NULL, &dmregular));
139   PetscCall(DMSetFromOptions(dmregular));
140   PetscCall(DMSetUp(dmregular));
141 
142   dx = 2.0 / ((PetscReal)mx);
143   PetscCall(DMDASetUniformCoordinates(dmregular, -1.0 + 0.5 * dx, 1.0 - 0.5 * dx, -1.0 + 0.5 * dx, 1.0 - 0.5 * dx, -1.0, 1.0));
144 
145   /* Create a DMShell for point location purposes */
146   PetscCall(DMShellCreate(PETSC_COMM_WORLD, &dmcell));
147   PetscCall(DMSetApplicationContext(dmcell, dmregular));
148   dmcell->ops->locatepoints = DMLocatePoints_DMDARegular;
149   dmcell->ops->getneighbors = DMGetNeighbors_DMDARegular;
150 
151   /* Create the swarm */
152   PetscCall(DMCreate(PETSC_COMM_WORLD, &dms));
153   PetscCall(DMSetType(dms, DMSWARM));
154   PetscCall(DMSetDimension(dms, 2));
155   PetscCall(PetscObjectSetName((PetscObject)dms, "Particles"));
156 
157   PetscCall(DMSwarmSetType(dms, DMSWARM_PIC));
158   PetscCall(DMSwarmSetCellDM(dms, dmcell));
159 
160   PetscCall(DMSwarmRegisterPetscDatatypeField(dms, "itag", 1, PETSC_INT));
161   PetscCall(DMSwarmFinalizeFieldRegister(dms));
162   {
163     PetscInt           si, sj, milocal, mjlocal;
164     const PetscScalar *LA_coors;
165     Vec                coors;
166     PetscInt           cnt;
167 
168     PetscCall(DMDAGetCorners(dmregular, &si, &sj, NULL, &milocal, &mjlocal, NULL));
169     PetscCall(DMGetCoordinates(dmregular, &coors));
170     PetscCall(VecGetArrayRead(coors, &LA_coors));
171     PetscCall(DMSwarmSetLocalSizes(dms, milocal * mjlocal, 4));
172     PetscCall(DMSwarmGetLocalSize(dms, &nlocal));
173     PetscCall(DMSwarmGetField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array));
174     cnt = 0;
175     PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &rand));
176     PetscCall(PetscRandomSetInterval(rand, -1.0, 1.0));
177     for (p = 0; p < nlocal; p++) {
178       PetscReal px, py, rx, ry, r2;
179 
180       PetscCall(PetscRandomGetValueReal(rand, &rx));
181       PetscCall(PetscRandomGetValueReal(rand, &ry));
182 
183       px = PetscRealPart(LA_coors[2 * p + 0]) + 0.1 * rx * dx;
184       py = PetscRealPart(LA_coors[2 * p + 1]) + 0.1 * ry * dx;
185 
186       r2 = px * px + py * py;
187       if (r2 < 0.75 * 0.75) {
188         array[bs * cnt + 0] = px;
189         array[bs * cnt + 1] = py;
190         cnt++;
191       }
192     }
193     PetscCall(PetscRandomDestroy(&rand));
194     PetscCall(DMSwarmRestoreField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array));
195     PetscCall(VecRestoreArrayRead(coors, &LA_coors));
196     PetscCall(DMSwarmSetLocalSizes(dms, cnt, 4));
197 
198     PetscCall(DMSwarmGetLocalSize(dms, &nlocal));
199     PetscCall(DMSwarmGetField(dms, "itag", &bs, NULL, (void **)&iarray));
200     for (p = 0; p < nlocal; p++) iarray[p] = (PetscInt)rank;
201     PetscCall(DMSwarmRestoreField(dms, "itag", &bs, NULL, (void **)&iarray));
202   }
203 
204   PetscCall(DMView(dms, PETSC_VIEWER_STDOUT_WORLD));
205   PetscCall(SwarmViewGP(dms, "step0"));
206 
207   dt = 0.1;
208   for (tk = 1; tk < 20; tk++) {
209     char prefix[PETSC_MAX_PATH_LEN];
210     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Step %" PetscInt_FMT " \n", tk));
211     /* push points */
212     PetscCall(DMSwarmGetLocalSize(dms, &nlocal));
213     PetscCall(DMSwarmGetField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array));
214     for (p = 0; p < nlocal; p++) {
215       PetscReal cx, cy, vx, vy;
216 
217       cx = array[2 * p];
218       cy = array[2 * p + 1];
219       vx = cy;
220       vy = -cx;
221 
222       array[2 * p] += dt * vx;
223       array[2 * p + 1] += dt * vy;
224     }
225     PetscCall(DMSwarmRestoreField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array));
226 
227     /* migrate points */
228     PetscCall(DMSwarmMigrate(dms, PETSC_TRUE));
229     /* view points */
230     PetscCall(PetscSNPrintf(prefix, PETSC_MAX_PATH_LEN - 1, "step%" PetscInt_FMT, tk));
231     /* should use the regular SwarmView() api, not one for a particular type */
232     PetscCall(SwarmViewGP(dms, prefix));
233   }
234   PetscCall(DMDestroy(&dmregular));
235   PetscCall(DMDestroy(&dmcell));
236   PetscCall(DMDestroy(&dms));
237   PetscFunctionReturn(0);
238 }
239 
240 int main(int argc, char **argv) {
241   PetscFunctionBeginUser;
242   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
243   PetscCall(ex3_1());
244   PetscCall(PetscFinalize());
245   return 0;
246 }
247 
248 /*TEST
249 
250    build:
251       requires: double !complex
252 
253    test:
254       filter: grep -v atomic
255       filter_output: grep -v atomic
256 TEST*/
257