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