Actual source code: swarmpic.c

  1: #define PETSCDM_DLL
  2: #include <petsc/private/dmswarmimpl.h>
  3: #include <petscsf.h>
  4: #include <petscdmda.h>
  5: #include <petscdmplex.h>
  6: #include <petscdt.h>
  7: #include "../src/dm/impls/swarm/data_bucket.h"

  9: #include <petsc/private/petscfeimpl.h>

 11: /* Coordinate insertition/addition API */
 12: /*@
 13:   DMSwarmSetPointsUniformCoordinates - Set point coordinates in a `DMSWARM` on a regular (ijk) grid

 15:   Collective

 17:   Input Parameters:
 18: + dm      - the `DMSWARM`
 19: . min     - minimum coordinate values in the x, y, z directions (array of length dim)
 20: . max     - maximum coordinate values in the x, y, z directions (array of length dim)
 21: . npoints - number of points in each spatial direction (array of length dim)
 22: - mode    - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)

 24:   Level: beginner

 26:   Notes:
 27:   When using mode = `INSERT_VALUES`, this method will reset the number of particles in the `DMSWARM`
 28:   to be `npoints[0]` x `npoints[1]` (2D) or `npoints[0]` x `npoints[1]` x `npoints[2]` (3D). When using mode = `ADD_VALUES`,
 29:   new points will be appended to any already existing in the `DMSWARM`

 31: .seealso: `DM`, `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
 32: @*/
 33: PetscErrorCode DMSwarmSetPointsUniformCoordinates(DM dm, PetscReal min[], PetscReal max[], PetscInt npoints[], InsertMode mode)
 34: {
 35:   PetscReal          lmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
 36:   PetscReal          lmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
 37:   PetscInt           i, j, k, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
 38:   const PetscScalar *_coor;
 39:   DM                 celldm;
 40:   PetscReal          dx[3];
 41:   PetscInt           _npoints[] = {0, 0, 1};
 42:   Vec                pos;
 43:   PetscScalar       *_pos;
 44:   PetscReal         *swarm_coor;
 45:   PetscInt          *swarm_cellid;
 46:   PetscSF            sfcell = NULL;
 47:   const PetscSFNode *LA_sfcell;

 49:   PetscFunctionBegin;
 50:   DMSWARMPICVALID(dm);
 51:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
 52:   PetscCall(DMGetLocalBoundingBox(celldm, lmin, lmax));
 53:   PetscCall(DMGetCoordinateDim(celldm, &bs));

 55:   for (b = 0; b < bs; b++) {
 56:     if (npoints[b] > 1) {
 57:       dx[b] = (max[b] - min[b]) / ((PetscReal)(npoints[b] - 1));
 58:     } else {
 59:       dx[b] = 0.0;
 60:     }
 61:     _npoints[b] = npoints[b];
 62:   }

 64:   /* determine number of points living in the bounding box */
 65:   n_estimate = 0;
 66:   for (k = 0; k < _npoints[2]; k++) {
 67:     for (j = 0; j < _npoints[1]; j++) {
 68:       for (i = 0; i < _npoints[0]; i++) {
 69:         PetscReal xp[] = {0.0, 0.0, 0.0};
 70:         PetscInt  ijk[3];
 71:         PetscBool point_inside = PETSC_TRUE;

 73:         ijk[0] = i;
 74:         ijk[1] = j;
 75:         ijk[2] = k;
 76:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
 77:         for (b = 0; b < bs; b++) {
 78:           if (xp[b] < lmin[b]) point_inside = PETSC_FALSE;
 79:           if (xp[b] > lmax[b]) point_inside = PETSC_FALSE;
 80:         }
 81:         if (point_inside) n_estimate++;
 82:       }
 83:     }
 84:   }

 86:   /* create candidate list */
 87:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
 88:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
 89:   PetscCall(VecSetBlockSize(pos, bs));
 90:   PetscCall(VecSetFromOptions(pos));
 91:   PetscCall(VecGetArray(pos, &_pos));

 93:   n_estimate = 0;
 94:   for (k = 0; k < _npoints[2]; k++) {
 95:     for (j = 0; j < _npoints[1]; j++) {
 96:       for (i = 0; i < _npoints[0]; i++) {
 97:         PetscReal xp[] = {0.0, 0.0, 0.0};
 98:         PetscInt  ijk[3];
 99:         PetscBool point_inside = PETSC_TRUE;

101:         ijk[0] = i;
102:         ijk[1] = j;
103:         ijk[2] = k;
104:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
105:         for (b = 0; b < bs; b++) {
106:           if (xp[b] < lmin[b]) point_inside = PETSC_FALSE;
107:           if (xp[b] > lmax[b]) point_inside = PETSC_FALSE;
108:         }
109:         if (point_inside) {
110:           for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = xp[b];
111:           n_estimate++;
112:         }
113:       }
114:     }
115:   }
116:   PetscCall(VecRestoreArray(pos, &_pos));

118:   /* locate points */
119:   PetscCall(DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell));
120:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
121:   n_found = 0;
122:   for (p = 0; p < n_estimate; p++) {
123:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
124:   }

126:   /* adjust size */
127:   if (mode == ADD_VALUES) {
128:     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
129:     n_new_est = n_curr + n_found;
130:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
131:   }
132:   if (mode == INSERT_VALUES) {
133:     n_curr    = 0;
134:     n_new_est = n_found;
135:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
136:   }

138:   /* initialize new coords, cell owners, pid */
139:   PetscCall(VecGetArrayRead(pos, &_coor));
140:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
141:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
142:   n_found = 0;
143:   for (p = 0; p < n_estimate; p++) {
144:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
145:       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
146:       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
147:       n_found++;
148:     }
149:   }
150:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
151:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
152:   PetscCall(VecRestoreArrayRead(pos, &_coor));

154:   PetscCall(PetscSFDestroy(&sfcell));
155:   PetscCall(VecDestroy(&pos));
156:   PetscFunctionReturn(PETSC_SUCCESS);
157: }

159: /*@
160:   DMSwarmSetPointCoordinates - Set point coordinates in a `DMSWARM` from a user defined list

162:   Collective

164:   Input Parameters:
165: + dm        - the `DMSWARM`
166: . npoints   - the number of points to insert
167: . coor      - the coordinate values
168: . redundant - if set to `PETSC_TRUE`, it is assumed that `npoints` and `coor` are only valid on rank 0 and should be broadcast to other ranks
169: - mode      - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)

171:   Level: beginner

173:   Notes:
174:   If the user has specified `redundant` as `PETSC_FALSE`, the cell `DM` will attempt to locate the coordinates provided by `coor` within
175:   its sub-domain. If they any values within `coor` are not located in the sub-domain, they will be ignored and will not get
176:   added to the `DMSWARM`.

178: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`, `DMSwarmSetPointsUniformCoordinates()`
179: @*/
180: PetscErrorCode DMSwarmSetPointCoordinates(DM dm, PetscInt npoints, PetscReal coor[], PetscBool redundant, InsertMode mode)
181: {
182:   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
183:   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
184:   PetscInt           i, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
185:   Vec                coorlocal;
186:   const PetscScalar *_coor;
187:   DM                 celldm;
188:   Vec                pos;
189:   PetscScalar       *_pos;
190:   PetscReal         *swarm_coor;
191:   PetscInt          *swarm_cellid;
192:   PetscSF            sfcell = NULL;
193:   const PetscSFNode *LA_sfcell;
194:   PetscReal         *my_coor;
195:   PetscInt           my_npoints;
196:   PetscMPIInt        rank;
197:   MPI_Comm           comm;

199:   PetscFunctionBegin;
200:   DMSWARMPICVALID(dm);
201:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
202:   PetscCallMPI(MPI_Comm_rank(comm, &rank));

204:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
205:   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
206:   PetscCall(VecGetSize(coorlocal, &N));
207:   PetscCall(VecGetBlockSize(coorlocal, &bs));
208:   N = N / bs;
209:   PetscCall(VecGetArrayRead(coorlocal, &_coor));
210:   for (i = 0; i < N; i++) {
211:     for (b = 0; b < bs; b++) {
212:       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
213:       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
214:     }
215:   }
216:   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));

218:   /* broadcast points from rank 0 if requested */
219:   if (redundant) {
220:     PetscMPIInt imy;

222:     my_npoints = npoints;
223:     PetscCallMPI(MPI_Bcast(&my_npoints, 1, MPIU_INT, 0, comm));

225:     if (rank > 0) { /* allocate space */
226:       PetscCall(PetscMalloc1(bs * my_npoints, &my_coor));
227:     } else {
228:       my_coor = coor;
229:     }
230:     PetscCall(PetscMPIIntCast(bs * my_npoints, &imy));
231:     PetscCallMPI(MPI_Bcast(my_coor, imy, MPIU_REAL, 0, comm));
232:   } else {
233:     my_npoints = npoints;
234:     my_coor    = coor;
235:   }

237:   /* determine the number of points living in the bounding box */
238:   n_estimate = 0;
239:   for (i = 0; i < my_npoints; i++) {
240:     PetscBool point_inside = PETSC_TRUE;

242:     for (b = 0; b < bs; b++) {
243:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
244:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
245:     }
246:     if (point_inside) n_estimate++;
247:   }

249:   /* create candidate list */
250:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
251:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
252:   PetscCall(VecSetBlockSize(pos, bs));
253:   PetscCall(VecSetFromOptions(pos));
254:   PetscCall(VecGetArray(pos, &_pos));

256:   n_estimate = 0;
257:   for (i = 0; i < my_npoints; i++) {
258:     PetscBool point_inside = PETSC_TRUE;

260:     for (b = 0; b < bs; b++) {
261:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
262:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
263:     }
264:     if (point_inside) {
265:       for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = my_coor[bs * i + b];
266:       n_estimate++;
267:     }
268:   }
269:   PetscCall(VecRestoreArray(pos, &_pos));

271:   /* locate points */
272:   PetscCall(DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell));

274:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
275:   n_found = 0;
276:   for (p = 0; p < n_estimate; p++) {
277:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
278:   }

280:   /* adjust size */
281:   if (mode == ADD_VALUES) {
282:     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
283:     n_new_est = n_curr + n_found;
284:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
285:   }
286:   if (mode == INSERT_VALUES) {
287:     n_curr    = 0;
288:     n_new_est = n_found;
289:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
290:   }

292:   /* initialize new coords, cell owners, pid */
293:   PetscCall(VecGetArrayRead(pos, &_coor));
294:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
295:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
296:   n_found = 0;
297:   for (p = 0; p < n_estimate; p++) {
298:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
299:       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
300:       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
301:       n_found++;
302:     }
303:   }
304:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
305:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
306:   PetscCall(VecRestoreArrayRead(pos, &_coor));

308:   if (redundant) {
309:     if (rank > 0) PetscCall(PetscFree(my_coor));
310:   }
311:   PetscCall(PetscSFDestroy(&sfcell));
312:   PetscCall(VecDestroy(&pos));
313:   PetscFunctionReturn(PETSC_SUCCESS);
314: }

316: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
317: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);

319: /*@
320:   DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell

322:   Not Collective

324:   Input Parameters:
325: + dm          - the `DMSWARM`
326: . layout_type - method used to fill each cell with the cell `DM`
327: - fill_param  - parameter controlling how many points per cell are added (the meaning of this parameter is dependent on the layout type)

329:   Level: beginner

331:   Notes:
332:   The insert method will reset any previous defined points within the `DMSWARM`.

334:   When using a `DMDA` both 2D and 3D are supported for all layout types provided you are using `DMDA_ELEMENT_Q1`.

336:   When using a `DMPLEX` the following case are supported\:
337: .vb
338:    (i) DMSWARMPIC_LAYOUT_REGULAR: 2D (triangle),
339:    (ii) DMSWARMPIC_LAYOUT_GAUSS: 2D and 3D provided the cell is a tri/tet or a quad/hex,
340:    (iii) DMSWARMPIC_LAYOUT_SUBDIVISION: 2D and 3D for quad/hex and 2D tri.
341: .ve

343: .seealso: `DMSWARM`, `DMSwarmPICLayoutType`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
344: @*/
345: PetscErrorCode DMSwarmInsertPointsUsingCellDM(DM dm, DMSwarmPICLayoutType layout_type, PetscInt fill_param)
346: {
347:   DM        celldm;
348:   PetscBool isDA, isPLEX;

350:   PetscFunctionBegin;
351:   DMSWARMPICVALID(dm);
352:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
353:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
354:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
355:   if (isDA) {
356:     PetscCall(private_DMSwarmInsertPointsUsingCellDM_DA(dm, celldm, layout_type, fill_param));
357:   } else if (isPLEX) {
358:     PetscCall(private_DMSwarmInsertPointsUsingCellDM_PLEX(dm, celldm, layout_type, fill_param));
359:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
360:   PetscFunctionReturn(PETSC_SUCCESS);
361: }

363: extern PetscErrorCode private_DMSwarmSetPointCoordinatesCellwise_PLEX(DM, DM, PetscInt, PetscReal *);

365: /*@C
366:   DMSwarmSetPointCoordinatesCellwise - Insert point coordinates (defined over the reference cell) within each cell

368:   Not Collective

370:   Input Parameters:
371: + dm      - the `DMSWARM`
372: . npoints - the number of points to insert in each cell
373: - xi      - the coordinates (defined in the local coordinate system for each cell) to insert

375:   Level: beginner

377:   Notes:
378:   The method will reset any previous defined points within the `DMSWARM`.
379:   Only supported for `DMPLEX`. If you are using a `DMDA` it is recommended to either use
380:   `DMSwarmInsertPointsUsingCellDM()`, or extract and set the coordinates yourself the following code
381: .vb
382:     PetscReal *coor;
383:     DMSwarmGetField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
384:     // user code to define the coordinates here
385:     DMSwarmRestoreField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
386: .ve

388: .seealso: `DMSWARM`, `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
389: @*/
390: PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
391: {
392:   DM        celldm;
393:   PetscBool isDA, isPLEX;

395:   PetscFunctionBegin;
396:   DMSWARMPICVALID(dm);
397:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
398:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
399:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
400:   PetscCheck(!isDA, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMPLEX. Recommended you use DMSwarmInsertPointsUsingCellDM()");
401:   if (isPLEX) {
402:     PetscCall(private_DMSwarmSetPointCoordinatesCellwise_PLEX(dm, celldm, npoints, xi));
403:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
404:   PetscFunctionReturn(PETSC_SUCCESS);
405: }

407: /*@
408:   DMSwarmCreatePointPerCellCount - Count the number of points within all cells in the cell DM

410:   Not Collective

412:   Input Parameter:
413: . dm - the `DMSWARM`

415:   Output Parameters:
416: + ncells - the number of cells in the cell `DM` (optional argument, pass `NULL` to ignore)
417: - count  - array of length ncells containing the number of points per cell

419:   Level: beginner

421:   Notes:
422:   The array count is allocated internally and must be free'd by the user.

424: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
425: @*/
426: PetscErrorCode DMSwarmCreatePointPerCellCount(DM dm, PetscInt *ncells, PetscInt **count)
427: {
428:   PetscBool isvalid;
429:   PetscInt  nel;
430:   PetscInt *sum;

432:   PetscFunctionBegin;
433:   PetscCall(DMSwarmSortGetIsValid(dm, &isvalid));
434:   nel = 0;
435:   if (isvalid) {
436:     PetscInt e;

438:     PetscCall(DMSwarmSortGetSizes(dm, &nel, NULL));

440:     PetscCall(PetscMalloc1(nel, &sum));
441:     for (e = 0; e < nel; e++) PetscCall(DMSwarmSortGetNumberOfPointsPerCell(dm, e, &sum[e]));
442:   } else {
443:     DM        celldm;
444:     PetscBool isda, isplex, isshell;
445:     PetscInt  p, npoints;
446:     PetscInt *swarm_cellid;

448:     /* get the number of cells */
449:     PetscCall(DMSwarmGetCellDM(dm, &celldm));
450:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isda));
451:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isplex));
452:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMSHELL, &isshell));
453:     if (isda) {
454:       PetscInt        _nel, _npe;
455:       const PetscInt *_element;

457:       PetscCall(DMDAGetElements(celldm, &_nel, &_npe, &_element));
458:       nel = _nel;
459:       PetscCall(DMDARestoreElements(celldm, &_nel, &_npe, &_element));
460:     } else if (isplex) {
461:       PetscInt ps, pe;

463:       PetscCall(DMPlexGetHeightStratum(celldm, 0, &ps, &pe));
464:       nel = pe - ps;
465:     } else if (isshell) {
466:       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);

468:       PetscCall(PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
469:       if (method_DMShellGetNumberOfCells) {
470:         PetscCall(method_DMShellGetNumberOfCells(celldm, &nel));
471:       } else
472:         SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for the DMSHELL object. User must provide a method via PetscObjectComposeFunction( (PetscObject)shelldm, \"DMGetNumberOfCells_C\", your_function_to_compute_number_of_cells);");
473:     } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");

475:     PetscCall(PetscMalloc1(nel, &sum));
476:     PetscCall(PetscArrayzero(sum, nel));
477:     PetscCall(DMSwarmGetLocalSize(dm, &npoints));
478:     PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
479:     for (p = 0; p < npoints; p++) {
480:       if (swarm_cellid[p] != DMLOCATEPOINT_POINT_NOT_FOUND) sum[swarm_cellid[p]]++;
481:     }
482:     PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
483:   }
484:   if (ncells) *ncells = nel;
485:   *count = sum;
486:   PetscFunctionReturn(PETSC_SUCCESS);
487: }

489: /*@
490:   DMSwarmGetNumSpecies - Get the number of particle species

492:   Not Collective

494:   Input Parameter:
495: . sw - the `DMSWARM`

497:   Output Parameters:
498: . Ns - the number of species

500:   Level: intermediate

502: .seealso: `DMSWARM`, `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
503: @*/
504: PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
505: {
506:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

508:   PetscFunctionBegin;
509:   *Ns = swarm->Ns;
510:   PetscFunctionReturn(PETSC_SUCCESS);
511: }

513: /*@
514:   DMSwarmSetNumSpecies - Set the number of particle species

516:   Not Collective

518:   Input Parameters:
519: + sw - the `DMSWARM`
520: - Ns - the number of species

522:   Level: intermediate

524: .seealso: `DMSWARM`, `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
525: @*/
526: PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
527: {
528:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

530:   PetscFunctionBegin;
531:   swarm->Ns = Ns;
532:   PetscFunctionReturn(PETSC_SUCCESS);
533: }

535: /*@C
536:   DMSwarmGetCoordinateFunction - Get the function setting initial particle positions, if it exists

538:   Not Collective

540:   Input Parameter:
541: . sw - the `DMSWARM`

543:   Output Parameter:
544: . coordFunc - the function setting initial particle positions, or `NULL`, see `PetscSimplePointFn` for the calling sequence

546:   Level: intermediate

548: .seealso: `DMSWARM`, `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`, `PetscSimplePointFn`
549: @*/
550: PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFn **coordFunc)
551: {
552:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

554:   PetscFunctionBegin;
556:   PetscAssertPointer(coordFunc, 2);
557:   *coordFunc = swarm->coordFunc;
558:   PetscFunctionReturn(PETSC_SUCCESS);
559: }

561: /*@C
562:   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions

564:   Not Collective

566:   Input Parameters:
567: + sw        - the `DMSWARM`
568: - coordFunc - the function setting initial particle positions, see `PetscSimplePointFn` for the calling sequence

570:   Level: intermediate

572: .seealso: `DMSWARM`, `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`, `PetscSimplePointFn`
573: @*/
574: PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFn *coordFunc)
575: {
576:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

578:   PetscFunctionBegin;
581:   swarm->coordFunc = coordFunc;
582:   PetscFunctionReturn(PETSC_SUCCESS);
583: }

585: /*@C
586:   DMSwarmGetVelocityFunction - Get the function setting initial particle velocities, if it exists

588:   Not Collective

590:   Input Parameter:
591: . sw - the `DMSWARM`

593:   Output Parameter:
594: . velFunc - the function setting initial particle velocities, or `NULL`, see `PetscSimplePointFn` for the calling sequence

596:   Level: intermediate

598: .seealso: `DMSWARM`, `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`, `PetscSimplePointFn`
599: @*/
600: PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFn **velFunc)
601: {
602:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

604:   PetscFunctionBegin;
606:   PetscAssertPointer(velFunc, 2);
607:   *velFunc = swarm->velFunc;
608:   PetscFunctionReturn(PETSC_SUCCESS);
609: }

611: /*@C
612:   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities

614:   Not Collective

616:   Input Parameters:
617: + sw      - the `DMSWARM`
618: - velFunc - the function setting initial particle velocities, see `PetscSimplePointFn` for the calling sequence

620:   Level: intermediate

622: .seealso: `DMSWARM`, `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`, `PetscSimplePointFn`
623: @*/
624: PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFn *velFunc)
625: {
626:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

628:   PetscFunctionBegin;
631:   swarm->velFunc = velFunc;
632:   PetscFunctionReturn(PETSC_SUCCESS);
633: }

635: /*@C
636:   DMSwarmComputeLocalSize - Compute the local number and distribution of particles based upon a density function

638:   Not Collective

640:   Input Parameters:
641: + sw      - The `DMSWARM`
642: . N       - The target number of particles
643: - density - The density field for the particle layout, normalized to unity

645:   Level: advanced

647:   Note:
648:   One particle will be created for each species.

650: .seealso: `DMSWARM`, `DMSwarmComputeLocalSizeFromOptions()`
651: @*/
652: PetscErrorCode DMSwarmComputeLocalSize(DM sw, PetscInt N, PetscProbFunc density)
653: {
654:   DM               dm;
655:   PetscQuadrature  quad;
656:   const PetscReal *xq, *wq;
657:   PetscReal       *n_int;
658:   PetscInt        *npc_s, *cellid, Ni;
659:   PetscReal        gmin[3], gmax[3], xi0[3];
660:   PetscInt         Ns, cStart, cEnd, c, dim, d, Nq, q, Np = 0, p, s;
661:   PetscBool        simplex;

663:   PetscFunctionBegin;
664:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
665:   PetscCall(DMSwarmGetCellDM(sw, &dm));
666:   PetscCall(DMGetDimension(dm, &dim));
667:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
668:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
669:   PetscCall(DMPlexIsSimplex(dm, &simplex));
670:   PetscCall(DMGetCoordinatesLocalSetUp(dm));
671:   if (simplex) PetscCall(PetscDTStroudConicalQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
672:   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
673:   PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, &xq, &wq));
674:   PetscCall(PetscCalloc2(Ns, &n_int, (cEnd - cStart) * Ns, &npc_s));
675:   /* Integrate the density function to get the number of particles in each cell */
676:   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
677:   for (c = 0; c < cEnd - cStart; ++c) {
678:     const PetscInt cell = c + cStart;
679:     PetscReal      v0[3], J[9], invJ[9], detJ, detJp = 2. / (gmax[0] - gmin[0]), xr[3], den;

681:     /* Have to transform quadrature points/weights to cell domain */
682:     PetscCall(DMPlexComputeCellGeometryFEM(dm, cell, NULL, v0, J, invJ, &detJ));
683:     PetscCall(PetscArrayzero(n_int, Ns));
684:     for (q = 0; q < Nq; ++q) {
685:       CoordinatesRefToReal(dim, dim, xi0, v0, J, &xq[q * dim], xr);
686:       /* Have to transform mesh to domain of definition of PDF, [-1, 1], and weight PDF by |J|/2 */
687:       xr[0] = detJp * (xr[0] - gmin[0]) - 1.;

689:       for (s = 0; s < Ns; ++s) {
690:         PetscCall(density(xr, NULL, &den));
691:         n_int[s] += (detJp * den) * (detJ * wq[q]) / (PetscReal)Ns;
692:       }
693:     }
694:     for (s = 0; s < Ns; ++s) {
695:       Ni = N;
696:       npc_s[c * Ns + s] += (PetscInt)(Ni * n_int[s]);
697:       Np += npc_s[c * Ns + s];
698:     }
699:   }
700:   PetscCall(PetscQuadratureDestroy(&quad));
701:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
702:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
703:   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
704:     for (s = 0; s < Ns; ++s) {
705:       for (q = 0; q < npc_s[c * Ns + s]; ++q, ++p) cellid[p] = c;
706:     }
707:   }
708:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
709:   PetscCall(PetscFree2(n_int, npc_s));
710:   PetscFunctionReturn(PETSC_SUCCESS);
711: }

713: /*@
714:   DMSwarmComputeLocalSizeFromOptions - Compute the local number and distribution of particles based upon a density function determined by options

716:   Not Collective

718:   Input Parameter:
719: . sw - The `DMSWARM`

721:   Level: advanced

723: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`
724: @*/
725: PetscErrorCode DMSwarmComputeLocalSizeFromOptions(DM sw)
726: {
727:   PetscProbFunc pdf;
728:   const char   *prefix;
729:   char          funcname[PETSC_MAX_PATH_LEN];
730:   PetscInt     *N, Ns, dim, n;
731:   PetscBool     flg;
732:   PetscMPIInt   size, rank;

734:   PetscFunctionBegin;
735:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
736:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
737:   PetscCall(PetscCalloc1(size, &N));
738:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
739:   n = size;
740:   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
741:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
742:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
743:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
744:   PetscCall(PetscOptionsString("-dm_swarm_coordinate_function", "Function to determine particle coordinates", "DMSwarmSetCoordinateFunction", funcname, funcname, sizeof(funcname), &flg));
745:   PetscOptionsEnd();
746:   if (flg) {
747:     PetscSimplePointFn *coordFunc;

749:     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
750:     PetscCall(PetscDLSym(NULL, funcname, (void **)&coordFunc));
751:     PetscCheck(coordFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
752:     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
753:     PetscCall(DMSwarmSetLocalSizes(sw, N[rank] * Ns, 0));
754:     PetscCall(DMSwarmSetCoordinateFunction(sw, coordFunc));
755:   } else {
756:     PetscCall(DMGetDimension(sw, &dim));
757:     PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
758:     PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_coordinate_density", &pdf, NULL, NULL));
759:     PetscCall(DMSwarmComputeLocalSize(sw, N[rank], pdf));
760:   }
761:   PetscCall(PetscFree(N));
762:   PetscFunctionReturn(PETSC_SUCCESS);
763: }

765: /*@
766:   DMSwarmInitializeCoordinates - Determine the initial coordinates of particles for a PIC method

768:   Not Collective

770:   Input Parameter:
771: . sw - The `DMSWARM`

773:   Level: advanced

775:   Note:
776:   Currently, we randomly place particles in their assigned cell

778: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeVelocities()`
779: @*/
780: PetscErrorCode DMSwarmInitializeCoordinates(DM sw)
781: {
782:   PetscSimplePointFn *coordFunc;
783:   PetscScalar        *weight;
784:   PetscReal          *x;
785:   PetscInt           *species;
786:   void               *ctx;
787:   PetscBool           removePoints = PETSC_TRUE;
788:   PetscDataType       dtype;
789:   PetscInt            Np, p, Ns, dim, d, bs;

791:   PetscFunctionBeginUser;
792:   PetscCall(DMGetDimension(sw, &dim));
793:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
794:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
795:   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));

797:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, &bs, &dtype, (void **)&x));
798:   PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight));
799:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
800:   if (coordFunc) {
801:     PetscCall(DMGetApplicationContext(sw, &ctx));
802:     for (p = 0; p < Np; ++p) {
803:       PetscScalar X[3];

805:       PetscCall((*coordFunc)(dim, 0., NULL, p, X, ctx));
806:       for (d = 0; d < dim; ++d) x[p * dim + d] = PetscRealPart(X[d]);
807:       weight[p]  = 1.0;
808:       species[p] = p % Ns;
809:     }
810:   } else {
811:     DM          dm;
812:     PetscRandom rnd;
813:     PetscReal   xi0[3];
814:     PetscInt    cStart, cEnd, c;

816:     PetscCall(DMSwarmGetCellDM(sw, &dm));
817:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
818:     PetscCall(DMGetApplicationContext(sw, &ctx));

820:     /* Set particle position randomly in cell, set weights to 1 */
821:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd));
822:     PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
823:     PetscCall(PetscRandomSetFromOptions(rnd));
824:     PetscCall(DMSwarmSortGetAccess(sw));
825:     for (d = 0; d < dim; ++d) xi0[d] = -1.0;
826:     for (c = cStart; c < cEnd; ++c) {
827:       PetscReal v0[3], J[9], invJ[9], detJ;
828:       PetscInt *pidx, Npc, q;

830:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
831:       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ));
832:       for (q = 0; q < Npc; ++q) {
833:         const PetscInt p = pidx[q];
834:         PetscReal      xref[3];

836:         for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &xref[d]));
837:         CoordinatesRefToReal(dim, dim, xi0, v0, J, xref, &x[p * dim]);

839:         weight[p]  = 1.0 / Np;
840:         species[p] = p % Ns;
841:       }
842:       PetscCall(PetscFree(pidx));
843:     }
844:     PetscCall(PetscRandomDestroy(&rnd));
845:     PetscCall(DMSwarmSortRestoreAccess(sw));
846:   }
847:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
848:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
849:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));

851:   PetscCall(DMSwarmMigrate(sw, removePoints));
852:   PetscCall(DMLocalizeCoordinates(sw));
853:   PetscFunctionReturn(PETSC_SUCCESS);
854: }

856: /*@C
857:   DMSwarmInitializeVelocities - Set the initial velocities of particles using a distribution.

859:   Collective

861:   Input Parameters:
862: + sw      - The `DMSWARM` object
863: . sampler - A function which uniformly samples the velocity PDF
864: - v0      - The velocity scale for nondimensionalization for each species

866:   Level: advanced

868:   Note:
869:   If `v0` is zero for the first species, all velocities are set to zero. If it is zero for any other species, the effect will be to give that species zero velocity.

871: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocitiesFromOptions()`
872: @*/
873: PetscErrorCode DMSwarmInitializeVelocities(DM sw, PetscProbFunc sampler, const PetscReal v0[])
874: {
875:   PetscSimplePointFn *velFunc;
876:   PetscReal          *v;
877:   PetscInt           *species;
878:   void               *ctx;
879:   PetscInt            dim, Np, p;

881:   PetscFunctionBegin;
882:   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));

884:   PetscCall(DMGetDimension(sw, &dim));
885:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
886:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
887:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
888:   if (v0[0] == 0.) {
889:     PetscCall(PetscArrayzero(v, Np * dim));
890:   } else if (velFunc) {
891:     PetscCall(DMGetApplicationContext(sw, &ctx));
892:     for (p = 0; p < Np; ++p) {
893:       PetscInt    s = species[p], d;
894:       PetscScalar vel[3];

896:       PetscCall((*velFunc)(dim, 0., NULL, p, vel, ctx));
897:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
898:     }
899:   } else {
900:     PetscRandom rnd;

902:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
903:     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
904:     PetscCall(PetscRandomSetFromOptions(rnd));

906:     for (p = 0; p < Np; ++p) {
907:       PetscInt  s = species[p], d;
908:       PetscReal a[3], vel[3];

910:       for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &a[d]));
911:       PetscCall(sampler(a, NULL, vel));
912:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * vel[d];
913:     }
914:     PetscCall(PetscRandomDestroy(&rnd));
915:   }
916:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
917:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
918:   PetscFunctionReturn(PETSC_SUCCESS);
919: }

921: /*@
922:   DMSwarmInitializeVelocitiesFromOptions - Set the initial velocities of particles using a distribution determined from options.

924:   Collective

926:   Input Parameters:
927: + sw - The `DMSWARM` object
928: - v0 - The velocity scale for nondimensionalization for each species

930:   Level: advanced

932: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
933: @*/
934: PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
935: {
936:   PetscProbFunc sampler;
937:   PetscInt      dim;
938:   const char   *prefix;
939:   char          funcname[PETSC_MAX_PATH_LEN];
940:   PetscBool     flg;

942:   PetscFunctionBegin;
943:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
944:   PetscCall(PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg));
945:   PetscOptionsEnd();
946:   if (flg) {
947:     PetscSimplePointFn *velFunc;

949:     PetscCall(PetscDLSym(NULL, funcname, (void **)&velFunc));
950:     PetscCheck(velFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
951:     PetscCall(DMSwarmSetVelocityFunction(sw, velFunc));
952:   }
953:   PetscCall(DMGetDimension(sw, &dim));
954:   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
955:   PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_velocity_density", NULL, NULL, &sampler));
956:   PetscCall(DMSwarmInitializeVelocities(sw, sampler, v0));
957:   PetscFunctionReturn(PETSC_SUCCESS);
958: }