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: /*@C
 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: PETSC_EXTERN 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: /*@C
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: PETSC_EXTERN 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:     my_npoints = npoints;
221:     PetscCallMPI(MPI_Bcast(&my_npoints, 1, MPIU_INT, 0, comm));

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

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

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

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

253:   n_estimate = 0;
254:   for (i = 0; i < my_npoints; i++) {
255:     PetscBool point_inside = PETSC_TRUE;

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

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

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

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

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

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

313: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
314: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);

316: /*@C
317:   DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell

319:   Not Collective

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

326:   Level: beginner

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

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

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

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

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

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

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

365:   Not Collective

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

372:   Level: beginner

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

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

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

404: /*@C
405:   DMSwarmCreatePointPerCellCount - Count the number of points within all cells in the cell DM

407:   Not Collective

409:   Input Parameter:
410: . dm - the `DMSWARM`

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

416:   Level: beginner

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

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

429:   PetscFunctionBegin;
430:   PetscCall(DMSwarmSortGetIsValid(dm, &isvalid));
431:   nel = 0;
432:   if (isvalid) {
433:     PetscInt e;

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

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

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

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

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

465:       PetscCall(PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
466:       if (method_DMShellGetNumberOfCells) {
467:         PetscCall(method_DMShellGetNumberOfCells(celldm, &nel));
468:       } else
469:         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);");
470:     } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");

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

486: /*@
487:   DMSwarmGetNumSpecies - Get the number of particle species

489:   Not Collective

491:   Input Parameter:
492: . sw - the `DMSWARM`

494:   Output Parameters:
495: . Ns - the number of species

497:   Level: intermediate

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

505:   PetscFunctionBegin;
506:   *Ns = swarm->Ns;
507:   PetscFunctionReturn(PETSC_SUCCESS);
508: }

510: /*@
511:   DMSwarmSetNumSpecies - Set the number of particle species

513:   Not Collective

515:   Input Parameters:
516: + sw - the `DMSWARM`
517: - Ns - the number of species

519:   Level: intermediate

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

527:   PetscFunctionBegin;
528:   swarm->Ns = Ns;
529:   PetscFunctionReturn(PETSC_SUCCESS);
530: }

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

535:   Not Collective

537:   Input Parameter:
538: . sw - the `DMSWARM`

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

543:   Level: intermediate

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

551:   PetscFunctionBegin;
553:   PetscAssertPointer(coordFunc, 2);
554:   *coordFunc = swarm->coordFunc;
555:   PetscFunctionReturn(PETSC_SUCCESS);
556: }

558: /*@C
559:   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions

561:   Not Collective

563:   Input Parameters:
564: + sw        - the `DMSWARM`
565: - coordFunc - the function setting initial particle positions, see `PetscSimplePointFn` for the calling sequence

567:   Level: intermediate

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

575:   PetscFunctionBegin;
578:   swarm->coordFunc = coordFunc;
579:   PetscFunctionReturn(PETSC_SUCCESS);
580: }

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

585:   Not Collective

587:   Input Parameter:
588: . sw - the `DMSWARM`

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

593:   Level: intermediate

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

601:   PetscFunctionBegin;
603:   PetscAssertPointer(velFunc, 2);
604:   *velFunc = swarm->velFunc;
605:   PetscFunctionReturn(PETSC_SUCCESS);
606: }

608: /*@C
609:   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities

611:   Not Collective

613:   Input Parameters:
614: + sw      - the `DMSWARM`
615: - velFunc - the function setting initial particle velocities, see `PetscSimplePointFn` for the calling sequence

617:   Level: intermediate

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

625:   PetscFunctionBegin;
628:   swarm->velFunc = velFunc;
629:   PetscFunctionReturn(PETSC_SUCCESS);
630: }

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

635:   Not Collective

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

642:   Level: advanced

644:   Note:
645:   One particle will be created for each species.

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

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

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

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

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

713:   Not Collective

715:   Input Parameter:
716: . sw - The `DMSWARM`

718:   Level: advanced

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

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

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

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

765:   Not Collective

767:   Input Parameter:
768: . sw - The `DMSWARM`

770:   Level: advanced

772:   Note:
773:   Currently, we randomly place particles in their assigned cell

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

788:   PetscFunctionBeginUser;
789:   PetscCall(DMGetDimension(sw, &dim));
790:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
791:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
792:   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));

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

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

813:     PetscCall(DMSwarmGetCellDM(sw, &dm));
814:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
815:     PetscCall(DMGetApplicationContext(sw, &ctx));

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

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

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

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

848:   PetscCall(DMSwarmMigrate(sw, removePoints));
849:   PetscCall(DMLocalizeCoordinates(sw));
850:   PetscFunctionReturn(PETSC_SUCCESS);
851: }

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

856:   Collective

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

863:   Level: advanced

865:   Note:
866:   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.

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

878:   PetscFunctionBegin;
879:   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));

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

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

899:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
900:     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
901:     PetscCall(PetscRandomSetFromOptions(rnd));

903:     for (p = 0; p < Np; ++p) {
904:       PetscInt  s = species[p], d;
905:       PetscReal a[3], vel[3];

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

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

921:   Collective

923:   Input Parameters:
924: + sw - The `DMSWARM` object
925: - v0 - The velocity scale for nondimensionalization for each species

927:   Level: advanced

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

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

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