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