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]*npoints[1] (2D) or npoints[0]*npoints[1]*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          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
 36:   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
 37:   PetscInt           i, j, k, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
 38:   Vec                coorlocal;
 39:   const PetscScalar *_coor;
 40:   DM                 celldm;
 41:   PetscReal          dx[3];
 42:   PetscInt           _npoints[] = {0, 0, 1};
 43:   Vec                pos;
 44:   PetscScalar       *_pos;
 45:   PetscReal         *swarm_coor;
 46:   PetscInt          *swarm_cellid;
 47:   PetscSF            sfcell = NULL;
 48:   const PetscSFNode *LA_sfcell;

 50:   PetscFunctionBegin;
 51:   DMSWARMPICVALID(dm);
 52:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
 53:   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
 54:   PetscCall(VecGetSize(coorlocal, &N));
 55:   PetscCall(VecGetBlockSize(coorlocal, &bs));
 56:   N = N / bs;
 57:   PetscCall(VecGetArrayRead(coorlocal, &_coor));
 58:   for (i = 0; i < N; i++) {
 59:     for (b = 0; b < bs; b++) {
 60:       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
 61:       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
 62:     }
 63:   }
 64:   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));

 66:   for (b = 0; b < bs; b++) {
 67:     if (npoints[b] > 1) {
 68:       dx[b] = (max[b] - min[b]) / ((PetscReal)(npoints[b] - 1));
 69:     } else {
 70:       dx[b] = 0.0;
 71:     }
 72:     _npoints[b] = npoints[b];
 73:   }

 75:   /* determine number of points living in the bounding box */
 76:   n_estimate = 0;
 77:   for (k = 0; k < _npoints[2]; k++) {
 78:     for (j = 0; j < _npoints[1]; j++) {
 79:       for (i = 0; i < _npoints[0]; i++) {
 80:         PetscReal xp[] = {0.0, 0.0, 0.0};
 81:         PetscInt  ijk[3];
 82:         PetscBool point_inside = PETSC_TRUE;

 84:         ijk[0] = i;
 85:         ijk[1] = j;
 86:         ijk[2] = k;
 87:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
 88:         for (b = 0; b < bs; b++) {
 89:           if (xp[b] < gmin[b]) point_inside = PETSC_FALSE;
 90:           if (xp[b] > gmax[b]) point_inside = PETSC_FALSE;
 91:         }
 92:         if (point_inside) n_estimate++;
 93:       }
 94:     }
 95:   }

 97:   /* create candidate list */
 98:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
 99:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
100:   PetscCall(VecSetBlockSize(pos, bs));
101:   PetscCall(VecSetFromOptions(pos));
102:   PetscCall(VecGetArray(pos, &_pos));

104:   n_estimate = 0;
105:   for (k = 0; k < _npoints[2]; k++) {
106:     for (j = 0; j < _npoints[1]; j++) {
107:       for (i = 0; i < _npoints[0]; i++) {
108:         PetscReal xp[] = {0.0, 0.0, 0.0};
109:         PetscInt  ijk[3];
110:         PetscBool point_inside = PETSC_TRUE;

112:         ijk[0] = i;
113:         ijk[1] = j;
114:         ijk[2] = k;
115:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
116:         for (b = 0; b < bs; b++) {
117:           if (xp[b] < gmin[b]) point_inside = PETSC_FALSE;
118:           if (xp[b] > gmax[b]) point_inside = PETSC_FALSE;
119:         }
120:         if (point_inside) {
121:           for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = xp[b];
122:           n_estimate++;
123:         }
124:       }
125:     }
126:   }
127:   PetscCall(VecRestoreArray(pos, &_pos));

129:   /* locate points */
130:   PetscCall(DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell));
131:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
132:   n_found = 0;
133:   for (p = 0; p < n_estimate; p++) {
134:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
135:   }

137:   /* adjust size */
138:   if (mode == ADD_VALUES) {
139:     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
140:     n_new_est = n_curr + n_found;
141:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
142:   }
143:   if (mode == INSERT_VALUES) {
144:     n_curr    = 0;
145:     n_new_est = n_found;
146:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
147:   }

149:   /* initialize new coords, cell owners, pid */
150:   PetscCall(VecGetArrayRead(pos, &_coor));
151:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
152:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
153:   n_found = 0;
154:   for (p = 0; p < n_estimate; p++) {
155:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
156:       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
157:       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
158:       n_found++;
159:     }
160:   }
161:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
162:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
163:   PetscCall(VecRestoreArrayRead(pos, &_coor));

165:   PetscCall(PetscSFDestroy(&sfcell));
166:   PetscCall(VecDestroy(&pos));
167:   PetscFunctionReturn(PETSC_SUCCESS);
168: }

170: /*@C
171:   DMSwarmSetPointCoordinates - Set point coordinates in a `DMSWARM` from a user defined list

173:   Collective

175:   Input Parameters:
176: + dm        - the `DMSWARM`
177: . npoints   - the number of points to insert
178: . coor      - the coordinate values
179: . 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
180: - mode      - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)

182:   Level: beginner

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

189: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`, `DMSwarmSetPointsUniformCoordinates()`
190: @*/
191: PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinates(DM dm, PetscInt npoints, PetscReal coor[], PetscBool redundant, InsertMode mode)
192: {
193:   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
194:   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
195:   PetscInt           i, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
196:   Vec                coorlocal;
197:   const PetscScalar *_coor;
198:   DM                 celldm;
199:   Vec                pos;
200:   PetscScalar       *_pos;
201:   PetscReal         *swarm_coor;
202:   PetscInt          *swarm_cellid;
203:   PetscSF            sfcell = NULL;
204:   const PetscSFNode *LA_sfcell;
205:   PetscReal         *my_coor;
206:   PetscInt           my_npoints;
207:   PetscMPIInt        rank;
208:   MPI_Comm           comm;

210:   PetscFunctionBegin;
211:   DMSWARMPICVALID(dm);
212:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
213:   PetscCallMPI(MPI_Comm_rank(comm, &rank));

215:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
216:   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
217:   PetscCall(VecGetSize(coorlocal, &N));
218:   PetscCall(VecGetBlockSize(coorlocal, &bs));
219:   N = N / bs;
220:   PetscCall(VecGetArrayRead(coorlocal, &_coor));
221:   for (i = 0; i < N; i++) {
222:     for (b = 0; b < bs; b++) {
223:       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
224:       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
225:     }
226:   }
227:   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));

229:   /* broadcast points from rank 0 if requested */
230:   if (redundant) {
231:     my_npoints = npoints;
232:     PetscCallMPI(MPI_Bcast(&my_npoints, 1, MPIU_INT, 0, comm));

234:     if (rank > 0) { /* allocate space */
235:       PetscCall(PetscMalloc1(bs * my_npoints, &my_coor));
236:     } else {
237:       my_coor = coor;
238:     }
239:     PetscCallMPI(MPI_Bcast(my_coor, bs * my_npoints, MPIU_REAL, 0, comm));
240:   } else {
241:     my_npoints = npoints;
242:     my_coor    = coor;
243:   }

245:   /* determine the number of points living in the bounding box */
246:   n_estimate = 0;
247:   for (i = 0; i < my_npoints; i++) {
248:     PetscBool point_inside = PETSC_TRUE;

250:     for (b = 0; b < bs; b++) {
251:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
252:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
253:     }
254:     if (point_inside) n_estimate++;
255:   }

257:   /* create candidate list */
258:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
259:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
260:   PetscCall(VecSetBlockSize(pos, bs));
261:   PetscCall(VecSetFromOptions(pos));
262:   PetscCall(VecGetArray(pos, &_pos));

264:   n_estimate = 0;
265:   for (i = 0; i < my_npoints; i++) {
266:     PetscBool point_inside = PETSC_TRUE;

268:     for (b = 0; b < bs; b++) {
269:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
270:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
271:     }
272:     if (point_inside) {
273:       for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = my_coor[bs * i + b];
274:       n_estimate++;
275:     }
276:   }
277:   PetscCall(VecRestoreArray(pos, &_pos));

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

282:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
283:   n_found = 0;
284:   for (p = 0; p < n_estimate; p++) {
285:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
286:   }

288:   /* adjust size */
289:   if (mode == ADD_VALUES) {
290:     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
291:     n_new_est = n_curr + n_found;
292:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
293:   }
294:   if (mode == INSERT_VALUES) {
295:     n_curr    = 0;
296:     n_new_est = n_found;
297:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
298:   }

300:   /* initialize new coords, cell owners, pid */
301:   PetscCall(VecGetArrayRead(pos, &_coor));
302:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
303:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
304:   n_found = 0;
305:   for (p = 0; p < n_estimate; p++) {
306:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
307:       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
308:       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
309:       n_found++;
310:     }
311:   }
312:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
313:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
314:   PetscCall(VecRestoreArrayRead(pos, &_coor));

316:   if (redundant) {
317:     if (rank > 0) PetscCall(PetscFree(my_coor));
318:   }
319:   PetscCall(PetscSFDestroy(&sfcell));
320:   PetscCall(VecDestroy(&pos));
321:   PetscFunctionReturn(PETSC_SUCCESS);
322: }

324: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
325: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);

327: /*@C
328:   DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell

330:   Not Collective

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

337:   Level: beginner

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

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

344:   When using a `DMPLEX` the following case are supported\:
345: .vb
346:    (i) DMSWARMPIC_LAYOUT_REGULAR: 2D (triangle),
347:    (ii) DMSWARMPIC_LAYOUT_GAUSS: 2D and 3D provided the cell is a tri/tet or a quad/hex,
348:    (iii) DMSWARMPIC_LAYOUT_SUBDIVISION: 2D and 3D for quad/hex and 2D tri.
349: .ve

351: .seealso: `DMSWARM`, `DMSwarmPICLayoutType`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
352: @*/
353: PETSC_EXTERN PetscErrorCode DMSwarmInsertPointsUsingCellDM(DM dm, DMSwarmPICLayoutType layout_type, PetscInt fill_param)
354: {
355:   DM        celldm;
356:   PetscBool isDA, isPLEX;

358:   PetscFunctionBegin;
359:   DMSWARMPICVALID(dm);
360:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
361:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
362:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
363:   if (isDA) {
364:     PetscCall(private_DMSwarmInsertPointsUsingCellDM_DA(dm, celldm, layout_type, fill_param));
365:   } else if (isPLEX) {
366:     PetscCall(private_DMSwarmInsertPointsUsingCellDM_PLEX(dm, celldm, layout_type, fill_param));
367:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
368:   PetscFunctionReturn(PETSC_SUCCESS);
369: }

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

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

376:   Not Collective

378:   Input Parameters:
379: + dm      - the `DMSWARM`
380: . npoints - the number of points to insert in each cell
381: - xi      - the coordinates (defined in the local coordinate system for each cell) to insert

383:   Level: beginner

385:   Notes:
386:   The method will reset any previous defined points within the `DMSWARM`.
387:   Only supported for `DMPLEX`. If you are using a `DMDA` it is recommended to either use
388:   `DMSwarmInsertPointsUsingCellDM()`, or extract and set the coordinates yourself the following code
389: .vb
390:     PetscReal *coor;
391:     DMSwarmGetField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
392:     // user code to define the coordinates here
393:     DMSwarmRestoreField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
394: .ve

396: .seealso: `DMSWARM`, `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
397: @*/
398: PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
399: {
400:   DM        celldm;
401:   PetscBool isDA, isPLEX;

403:   PetscFunctionBegin;
404:   DMSWARMPICVALID(dm);
405:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
406:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
407:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
408:   PetscCheck(!isDA, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMPLEX. Recommended you use DMSwarmInsertPointsUsingCellDM()");
409:   if (isPLEX) {
410:     PetscCall(private_DMSwarmSetPointCoordinatesCellwise_PLEX(dm, celldm, npoints, xi));
411:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
412:   PetscFunctionReturn(PETSC_SUCCESS);
413: }

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

418:   Not Collective

420:   Input Parameter:
421: . dm - the `DMSWARM`

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

427:   Level: beginner

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

432: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
433: @*/
434: PETSC_EXTERN PetscErrorCode DMSwarmCreatePointPerCellCount(DM dm, PetscInt *ncells, PetscInt **count)
435: {
436:   PetscBool isvalid;
437:   PetscInt  nel;
438:   PetscInt *sum;

440:   PetscFunctionBegin;
441:   PetscCall(DMSwarmSortGetIsValid(dm, &isvalid));
442:   nel = 0;
443:   if (isvalid) {
444:     PetscInt e;

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

448:     PetscCall(PetscMalloc1(nel, &sum));
449:     for (e = 0; e < nel; e++) PetscCall(DMSwarmSortGetNumberOfPointsPerCell(dm, e, &sum[e]));
450:   } else {
451:     DM        celldm;
452:     PetscBool isda, isplex, isshell;
453:     PetscInt  p, npoints;
454:     PetscInt *swarm_cellid;

456:     /* get the number of cells */
457:     PetscCall(DMSwarmGetCellDM(dm, &celldm));
458:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isda));
459:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isplex));
460:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMSHELL, &isshell));
461:     if (isda) {
462:       PetscInt        _nel, _npe;
463:       const PetscInt *_element;

465:       PetscCall(DMDAGetElements(celldm, &_nel, &_npe, &_element));
466:       nel = _nel;
467:       PetscCall(DMDARestoreElements(celldm, &_nel, &_npe, &_element));
468:     } else if (isplex) {
469:       PetscInt ps, pe;

471:       PetscCall(DMPlexGetHeightStratum(celldm, 0, &ps, &pe));
472:       nel = pe - ps;
473:     } else if (isshell) {
474:       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);

476:       PetscCall(PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
477:       if (method_DMShellGetNumberOfCells) {
478:         PetscCall(method_DMShellGetNumberOfCells(celldm, &nel));
479:       } else
480:         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);");
481:     } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");

483:     PetscCall(PetscMalloc1(nel, &sum));
484:     PetscCall(PetscArrayzero(sum, nel));
485:     PetscCall(DMSwarmGetLocalSize(dm, &npoints));
486:     PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
487:     for (p = 0; p < npoints; p++) {
488:       if (swarm_cellid[p] != DMLOCATEPOINT_POINT_NOT_FOUND) sum[swarm_cellid[p]]++;
489:     }
490:     PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
491:   }
492:   if (ncells) *ncells = nel;
493:   *count = sum;
494:   PetscFunctionReturn(PETSC_SUCCESS);
495: }

497: /*@
498:   DMSwarmGetNumSpecies - Get the number of particle species

500:   Not Collective

502:   Input Parameter:
503: . sw - the `DMSWARM`

505:   Output Parameters:
506: . Ns - the number of species

508:   Level: intermediate

510: .seealso: `DMSWARM`, `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
511: @*/
512: PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
513: {
514:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

516:   PetscFunctionBegin;
517:   *Ns = swarm->Ns;
518:   PetscFunctionReturn(PETSC_SUCCESS);
519: }

521: /*@
522:   DMSwarmSetNumSpecies - Set the number of particle species

524:   Not Collective

526:   Input Parameters:
527: + sw - the `DMSWARM`
528: - Ns - the number of species

530:   Level: intermediate

532: .seealso: `DMSWARM`, `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
533: @*/
534: PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
535: {
536:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

538:   PetscFunctionBegin;
539:   swarm->Ns = Ns;
540:   PetscFunctionReturn(PETSC_SUCCESS);
541: }

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

546:   Not Collective

548:   Input Parameter:
549: . sw - the `DMSWARM`

551:   Output Parameter:
552: . coordFunc - the function setting initial particle positions, or `NULL`

554:   Level: intermediate

556: .seealso: `DMSWARM`, `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
557: @*/
558: PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFunc *coordFunc)
559: {
560:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

562:   PetscFunctionBegin;
564:   PetscAssertPointer(coordFunc, 2);
565:   *coordFunc = swarm->coordFunc;
566:   PetscFunctionReturn(PETSC_SUCCESS);
567: }

569: /*@C
570:   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions

572:   Not Collective

574:   Input Parameters:
575: + sw        - the `DMSWARM`
576: - coordFunc - the function setting initial particle positions

578:   Level: intermediate

580: .seealso: `DMSWARM`, `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
581: @*/
582: PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFunc coordFunc)
583: {
584:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

586:   PetscFunctionBegin;
589:   swarm->coordFunc = coordFunc;
590:   PetscFunctionReturn(PETSC_SUCCESS);
591: }

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

596:   Not Collective

598:   Input Parameter:
599: . sw - the `DMSWARM`

601:   Output Parameter:
602: . velFunc - the function setting initial particle velocities, or `NULL`

604:   Level: intermediate

606: .seealso: `DMSWARM`, `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
607: @*/
608: PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFunc *velFunc)
609: {
610:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

612:   PetscFunctionBegin;
614:   PetscAssertPointer(velFunc, 2);
615:   *velFunc = swarm->velFunc;
616:   PetscFunctionReturn(PETSC_SUCCESS);
617: }

619: /*@C
620:   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities

622:   Not Collective

624:   Input Parameters:
625: + sw      - the `DMSWARM`
626: - velFunc - the function setting initial particle velocities

628:   Level: intermediate

630: .seealso: `DMSWARM`, `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
631: @*/
632: PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFunc velFunc)
633: {
634:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

636:   PetscFunctionBegin;
639:   swarm->velFunc = velFunc;
640:   PetscFunctionReturn(PETSC_SUCCESS);
641: }

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

646:   Not Collective

648:   Input Parameters:
649: + sw      - The `DMSWARM`
650: . N       - The target number of particles
651: - density - The density field for the particle layout, normalized to unity

653:   Level: advanced

655:   Note:
656:   One particle will be created for each species.

658: .seealso: `DMSWARM`, `DMSwarmComputeLocalSizeFromOptions()`
659: @*/
660: PetscErrorCode DMSwarmComputeLocalSize(DM sw, PetscInt N, PetscProbFunc density)
661: {
662:   DM               dm;
663:   PetscQuadrature  quad;
664:   const PetscReal *xq, *wq;
665:   PetscReal       *n_int;
666:   PetscInt        *npc_s, *cellid, Ni;
667:   PetscReal        gmin[3], gmax[3], xi0[3];
668:   PetscInt         Ns, cStart, cEnd, c, dim, d, Nq, q, Np = 0, p, s;
669:   PetscBool        simplex;

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

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

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

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

724:   Not Collective

726:   Input Parameter:
727: . sw - The `DMSWARM`

729:   Level: advanced

731: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`
732: @*/
733: PetscErrorCode DMSwarmComputeLocalSizeFromOptions(DM sw)
734: {
735:   PetscProbFunc pdf;
736:   const char   *prefix;
737:   char          funcname[PETSC_MAX_PATH_LEN];
738:   PetscInt     *N, Ns, dim, n;
739:   PetscBool     flg;
740:   PetscMPIInt   size, rank;

742:   PetscFunctionBegin;
743:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
744:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
745:   PetscCall(PetscCalloc1(size, &N));
746:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
747:   n = size;
748:   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
749:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
750:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
751:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
752:   PetscCall(PetscOptionsString("-dm_swarm_coordinate_function", "Function to determine particle coordinates", "DMSwarmSetCoordinateFunction", funcname, funcname, sizeof(funcname), &flg));
753:   PetscOptionsEnd();
754:   if (flg) {
755:     PetscSimplePointFunc coordFunc;

757:     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
758:     PetscCall(PetscDLSym(NULL, funcname, (void **)&coordFunc));
759:     PetscCheck(coordFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
760:     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
761:     PetscCall(DMSwarmSetLocalSizes(sw, N[rank] * Ns, 0));
762:     PetscCall(DMSwarmSetCoordinateFunction(sw, coordFunc));
763:   } else {
764:     PetscCall(DMGetDimension(sw, &dim));
765:     PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
766:     PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_coordinate_density", &pdf, NULL, NULL));
767:     PetscCall(DMSwarmComputeLocalSize(sw, N[rank], pdf));
768:   }
769:   PetscCall(PetscFree(N));
770:   PetscFunctionReturn(PETSC_SUCCESS);
771: }

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

776:   Not Collective

778:   Input Parameter:
779: . sw - The `DMSWARM`

781:   Level: advanced

783:   Note:
784:   Currently, we randomly place particles in their assigned cell

786: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeVelocities()`
787: @*/
788: PetscErrorCode DMSwarmInitializeCoordinates(DM sw)
789: {
790:   PetscSimplePointFunc coordFunc;
791:   PetscScalar         *weight;
792:   PetscReal           *x;
793:   PetscInt            *species;
794:   void                *ctx;
795:   PetscBool            removePoints = PETSC_TRUE;
796:   PetscDataType        dtype;
797:   PetscInt             Np, p, Ns, dim, d, bs;

799:   PetscFunctionBeginUser;
800:   PetscCall(DMGetDimension(sw, &dim));
801:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
802:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
803:   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));

805:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, &bs, &dtype, (void **)&x));
806:   PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight));
807:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
808:   if (coordFunc) {
809:     PetscCall(DMGetApplicationContext(sw, &ctx));
810:     for (p = 0; p < Np; ++p) {
811:       PetscScalar X[3];

813:       PetscCall((*coordFunc)(dim, 0., NULL, p, X, ctx));
814:       for (d = 0; d < dim; ++d) x[p * dim + d] = PetscRealPart(X[d]);
815:       weight[p]  = 1.0;
816:       species[p] = p % Ns;
817:     }
818:   } else {
819:     DM          dm;
820:     PetscRandom rnd;
821:     PetscReal   xi0[3];
822:     PetscInt    cStart, cEnd, c;

824:     PetscCall(DMSwarmGetCellDM(sw, &dm));
825:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
826:     PetscCall(DMGetApplicationContext(sw, &ctx));

828:     /* Set particle position randomly in cell, set weights to 1 */
829:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd));
830:     PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
831:     PetscCall(PetscRandomSetFromOptions(rnd));
832:     PetscCall(DMSwarmSortGetAccess(sw));
833:     for (d = 0; d < dim; ++d) xi0[d] = -1.0;
834:     for (c = cStart; c < cEnd; ++c) {
835:       PetscReal v0[3], J[9], invJ[9], detJ;
836:       PetscInt *pidx, Npc, q;

838:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
839:       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ));
840:       for (q = 0; q < Npc; ++q) {
841:         const PetscInt p = pidx[q];
842:         PetscReal      xref[3];

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

847:         weight[p]  = 1.0 / Np;
848:         species[p] = p % Ns;
849:       }
850:       PetscCall(PetscFree(pidx));
851:     }
852:     PetscCall(PetscRandomDestroy(&rnd));
853:     PetscCall(DMSwarmSortRestoreAccess(sw));
854:   }
855:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
856:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
857:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));

859:   PetscCall(DMSwarmMigrate(sw, removePoints));
860:   PetscCall(DMLocalizeCoordinates(sw));
861:   PetscFunctionReturn(PETSC_SUCCESS);
862: }

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

867:   Collective

869:   Input Parameters:
870: + sw      - The `DMSWARM` object
871: . sampler - A function which uniformly samples the velocity PDF
872: - v0      - The velocity scale for nondimensionalization for each species

874:   Level: advanced

876:   Note:
877:   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.

879: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocitiesFromOptions()`
880: @*/
881: PetscErrorCode DMSwarmInitializeVelocities(DM sw, PetscProbFunc sampler, const PetscReal v0[])
882: {
883:   PetscSimplePointFunc velFunc;
884:   PetscReal           *v;
885:   PetscInt            *species;
886:   void                *ctx;
887:   PetscInt             dim, Np, p;

889:   PetscFunctionBegin;
890:   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));

892:   PetscCall(DMGetDimension(sw, &dim));
893:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
894:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
895:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
896:   if (v0[0] == 0.) {
897:     PetscCall(PetscArrayzero(v, Np * dim));
898:   } else if (velFunc) {
899:     PetscCall(DMGetApplicationContext(sw, &ctx));
900:     for (p = 0; p < Np; ++p) {
901:       PetscInt    s = species[p], d;
902:       PetscScalar vel[3];

904:       PetscCall((*velFunc)(dim, 0., NULL, p, vel, ctx));
905:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
906:     }
907:   } else {
908:     PetscRandom rnd;

910:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
911:     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
912:     PetscCall(PetscRandomSetFromOptions(rnd));

914:     for (p = 0; p < Np; ++p) {
915:       PetscInt  s = species[p], d;
916:       PetscReal a[3], vel[3];

918:       for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &a[d]));
919:       PetscCall(sampler(a, NULL, vel));
920:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * vel[d];
921:     }
922:     PetscCall(PetscRandomDestroy(&rnd));
923:   }
924:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
925:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
926:   PetscFunctionReturn(PETSC_SUCCESS);
927: }

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

932:   Collective

934:   Input Parameters:
935: + sw - The `DMSWARM` object
936: - v0 - The velocity scale for nondimensionalization for each species

938:   Level: advanced

940: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
941: @*/
942: PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
943: {
944:   PetscProbFunc sampler;
945:   PetscInt      dim;
946:   const char   *prefix;
947:   char          funcname[PETSC_MAX_PATH_LEN];
948:   PetscBool     flg;

950:   PetscFunctionBegin;
951:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
952:   PetscCall(PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg));
953:   PetscOptionsEnd();
954:   if (flg) {
955:     PetscSimplePointFunc velFunc;

957:     PetscCall(PetscDLSym(NULL, funcname, (void **)&velFunc));
958:     PetscCheck(velFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
959:     PetscCall(DMSwarmSetVelocityFunction(sw, velFunc));
960:   }
961:   PetscCall(DMGetDimension(sw, &dim));
962:   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
963:   PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_velocity_density", NULL, NULL, &sampler));
964:   PetscCall(DMSwarmInitializeVelocities(sw, sampler, v0));
965:   PetscFunctionReturn(PETSC_SUCCESS);
966: }