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: /*
 12:  Error checking to ensure the swarm type is correct and that a cell DM has been set
 13: */
 14: #define DMSWARMPICVALID(dm) \
 15:   do { \
 16:     DM_Swarm *_swarm = (DM_Swarm *)(dm)->data; \
 17:     PetscCheck(_swarm->swarm_type == DMSWARM_PIC, PetscObjectComm((PetscObject)(dm)), PETSC_ERR_SUP, "Valid only for DMSwarm-PIC. You must call DMSwarmSetType(dm,DMSWARM_PIC)"); \
 18:     PetscCheck(_swarm->dmcell, PetscObjectComm((PetscObject)(dm)), PETSC_ERR_SUP, "Valid only for DMSwarmPIC if the cell DM is set. You must call DMSwarmSetCellDM(dm,celldm)"); \
 19:   } while (0)

 21: /* Coordinate insertition/addition API */
 22: /*@C
 23:    DMSwarmSetPointsUniformCoordinates - Set point coordinates in a `DMSWARM` on a regular (ijk) grid

 25:    Collective

 27:    Input parameters:
 28: +  dm - the `DMSWARM`
 29: .  min - minimum coordinate values in the x, y, z directions (array of length dim)
 30: .  max - maximum coordinate values in the x, y, z directions (array of length dim)
 31: .  npoints - number of points in each spatial direction (array of length dim)
 32: -  mode - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)

 34:    Level: beginner

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

 41: .seealso: `DM`, `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
 42: @*/
 43: PETSC_EXTERN PetscErrorCode DMSwarmSetPointsUniformCoordinates(DM dm, PetscReal min[], PetscReal max[], PetscInt npoints[], InsertMode mode)
 44: {
 45:   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
 46:   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
 47:   PetscInt           i, j, k, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
 48:   Vec                coorlocal;
 49:   const PetscScalar *_coor;
 50:   DM                 celldm;
 51:   PetscReal          dx[3];
 52:   PetscInt           _npoints[] = {0, 0, 1};
 53:   Vec                pos;
 54:   PetscScalar       *_pos;
 55:   PetscReal         *swarm_coor;
 56:   PetscInt          *swarm_cellid;
 57:   PetscSF            sfcell = NULL;
 58:   const PetscSFNode *LA_sfcell;

 60:   PetscFunctionBegin;
 61:   DMSWARMPICVALID(dm);
 62:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
 63:   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
 64:   PetscCall(VecGetSize(coorlocal, &N));
 65:   PetscCall(VecGetBlockSize(coorlocal, &bs));
 66:   N = N / bs;
 67:   PetscCall(VecGetArrayRead(coorlocal, &_coor));
 68:   for (i = 0; i < N; i++) {
 69:     for (b = 0; b < bs; b++) {
 70:       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
 71:       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
 72:     }
 73:   }
 74:   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));

 76:   for (b = 0; b < bs; b++) {
 77:     if (npoints[b] > 1) {
 78:       dx[b] = (max[b] - min[b]) / ((PetscReal)(npoints[b] - 1));
 79:     } else {
 80:       dx[b] = 0.0;
 81:     }
 82:     _npoints[b] = npoints[b];
 83:   }

 85:   /* determine number of points living in the bounding box */
 86:   n_estimate = 0;
 87:   for (k = 0; k < _npoints[2]; k++) {
 88:     for (j = 0; j < _npoints[1]; j++) {
 89:       for (i = 0; i < _npoints[0]; i++) {
 90:         PetscReal xp[] = {0.0, 0.0, 0.0};
 91:         PetscInt  ijk[3];
 92:         PetscBool point_inside = PETSC_TRUE;

 94:         ijk[0] = i;
 95:         ijk[1] = j;
 96:         ijk[2] = k;
 97:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
 98:         for (b = 0; b < bs; b++) {
 99:           if (xp[b] < gmin[b]) point_inside = PETSC_FALSE;
100:           if (xp[b] > gmax[b]) point_inside = PETSC_FALSE;
101:         }
102:         if (point_inside) n_estimate++;
103:       }
104:     }
105:   }

107:   /* create candidate list */
108:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
109:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
110:   PetscCall(VecSetBlockSize(pos, bs));
111:   PetscCall(VecSetFromOptions(pos));
112:   PetscCall(VecGetArray(pos, &_pos));

114:   n_estimate = 0;
115:   for (k = 0; k < _npoints[2]; k++) {
116:     for (j = 0; j < _npoints[1]; j++) {
117:       for (i = 0; i < _npoints[0]; i++) {
118:         PetscReal xp[] = {0.0, 0.0, 0.0};
119:         PetscInt  ijk[3];
120:         PetscBool point_inside = PETSC_TRUE;

122:         ijk[0] = i;
123:         ijk[1] = j;
124:         ijk[2] = k;
125:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
126:         for (b = 0; b < bs; b++) {
127:           if (xp[b] < gmin[b]) point_inside = PETSC_FALSE;
128:           if (xp[b] > gmax[b]) point_inside = PETSC_FALSE;
129:         }
130:         if (point_inside) {
131:           for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = xp[b];
132:           n_estimate++;
133:         }
134:       }
135:     }
136:   }
137:   PetscCall(VecRestoreArray(pos, &_pos));

139:   /* locate points */
140:   PetscCall(DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell));
141:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
142:   n_found = 0;
143:   for (p = 0; p < n_estimate; p++) {
144:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
145:   }

147:   /* adjust size */
148:   if (mode == ADD_VALUES) {
149:     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
150:     n_new_est = n_curr + n_found;
151:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
152:   }
153:   if (mode == INSERT_VALUES) {
154:     n_curr    = 0;
155:     n_new_est = n_found;
156:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
157:   }

159:   /* initialize new coords, cell owners, pid */
160:   PetscCall(VecGetArrayRead(pos, &_coor));
161:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
162:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
163:   n_found = 0;
164:   for (p = 0; p < n_estimate; p++) {
165:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
166:       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
167:       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
168:       n_found++;
169:     }
170:   }
171:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
172:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
173:   PetscCall(VecRestoreArrayRead(pos, &_coor));

175:   PetscCall(PetscSFDestroy(&sfcell));
176:   PetscCall(VecDestroy(&pos));
177:   PetscFunctionReturn(PETSC_SUCCESS);
178: }

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

183:    Collective

185:    Input parameters:
186: +  dm - the `DMSWARM`
187: .  npoints - the number of points to insert
188: .  coor - the coordinate values
189: .  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
190: -  mode - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)

192:    Level: beginner

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

199: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`, `DMSwarmSetPointsUniformCoordinates()`
200: @*/
201: PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinates(DM dm, PetscInt npoints, PetscReal coor[], PetscBool redundant, InsertMode mode)
202: {
203:   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
204:   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
205:   PetscInt           i, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
206:   Vec                coorlocal;
207:   const PetscScalar *_coor;
208:   DM                 celldm;
209:   Vec                pos;
210:   PetscScalar       *_pos;
211:   PetscReal         *swarm_coor;
212:   PetscInt          *swarm_cellid;
213:   PetscSF            sfcell = NULL;
214:   const PetscSFNode *LA_sfcell;
215:   PetscReal         *my_coor;
216:   PetscInt           my_npoints;
217:   PetscMPIInt        rank;
218:   MPI_Comm           comm;

220:   PetscFunctionBegin;
221:   DMSWARMPICVALID(dm);
222:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
223:   PetscCallMPI(MPI_Comm_rank(comm, &rank));

225:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
226:   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
227:   PetscCall(VecGetSize(coorlocal, &N));
228:   PetscCall(VecGetBlockSize(coorlocal, &bs));
229:   N = N / bs;
230:   PetscCall(VecGetArrayRead(coorlocal, &_coor));
231:   for (i = 0; i < N; i++) {
232:     for (b = 0; b < bs; b++) {
233:       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
234:       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
235:     }
236:   }
237:   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));

239:   /* broadcast points from rank 0 if requested */
240:   if (redundant) {
241:     my_npoints = npoints;
242:     PetscCallMPI(MPI_Bcast(&my_npoints, 1, MPIU_INT, 0, comm));

244:     if (rank > 0) { /* allocate space */
245:       PetscCall(PetscMalloc1(bs * my_npoints, &my_coor));
246:     } else {
247:       my_coor = coor;
248:     }
249:     PetscCallMPI(MPI_Bcast(my_coor, bs * my_npoints, MPIU_REAL, 0, comm));
250:   } else {
251:     my_npoints = npoints;
252:     my_coor    = coor;
253:   }

255:   /* determine the number of points living in the bounding box */
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) n_estimate++;
265:   }

267:   /* create candidate list */
268:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
269:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
270:   PetscCall(VecSetBlockSize(pos, bs));
271:   PetscCall(VecSetFromOptions(pos));
272:   PetscCall(VecGetArray(pos, &_pos));

274:   n_estimate = 0;
275:   for (i = 0; i < my_npoints; i++) {
276:     PetscBool point_inside = PETSC_TRUE;

278:     for (b = 0; b < bs; b++) {
279:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
280:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
281:     }
282:     if (point_inside) {
283:       for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = my_coor[bs * i + b];
284:       n_estimate++;
285:     }
286:   }
287:   PetscCall(VecRestoreArray(pos, &_pos));

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

292:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
293:   n_found = 0;
294:   for (p = 0; p < n_estimate; p++) {
295:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
296:   }

298:   /* adjust size */
299:   if (mode == ADD_VALUES) {
300:     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
301:     n_new_est = n_curr + n_found;
302:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
303:   }
304:   if (mode == INSERT_VALUES) {
305:     n_curr    = 0;
306:     n_new_est = n_found;
307:     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
308:   }

310:   /* initialize new coords, cell owners, pid */
311:   PetscCall(VecGetArrayRead(pos, &_coor));
312:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
313:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
314:   n_found = 0;
315:   for (p = 0; p < n_estimate; p++) {
316:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
317:       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
318:       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
319:       n_found++;
320:     }
321:   }
322:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
323:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
324:   PetscCall(VecRestoreArrayRead(pos, &_coor));

326:   if (redundant) {
327:     if (rank > 0) PetscCall(PetscFree(my_coor));
328:   }
329:   PetscCall(PetscSFDestroy(&sfcell));
330:   PetscCall(VecDestroy(&pos));
331:   PetscFunctionReturn(PETSC_SUCCESS);
332: }

334: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
335: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);

337: /*@C
338:    DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell

340:    Not Collective

342:    Input parameters:
343: +  dm - the `DMSWARM`
344: .  layout_type - method used to fill each cell with the cell `DM`
345: -  fill_param - parameter controlling how many points per cell are added (the meaning of this parameter is dependent on the layout type)

347:    Level: beginner

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

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

354:    When using a `DMPLEX` the following case are supported:
355: .vb
356:    (i) DMSWARMPIC_LAYOUT_REGULAR: 2D (triangle),
357:    (ii) DMSWARMPIC_LAYOUT_GAUSS: 2D and 3D provided the cell is a tri/tet or a quad/hex,
358:    (iii) DMSWARMPIC_LAYOUT_SUBDIVISION: 2D and 3D for quad/hex and 2D tri.
359: .ve

361: .seealso: `DMSWARM`, `DMSwarmPICLayoutType`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
362: @*/
363: PETSC_EXTERN PetscErrorCode DMSwarmInsertPointsUsingCellDM(DM dm, DMSwarmPICLayoutType layout_type, PetscInt fill_param)
364: {
365:   DM        celldm;
366:   PetscBool isDA, isPLEX;

368:   PetscFunctionBegin;
369:   DMSWARMPICVALID(dm);
370:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
371:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
372:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
373:   if (isDA) {
374:     PetscCall(private_DMSwarmInsertPointsUsingCellDM_DA(dm, celldm, layout_type, fill_param));
375:   } else if (isPLEX) {
376:     PetscCall(private_DMSwarmInsertPointsUsingCellDM_PLEX(dm, celldm, layout_type, fill_param));
377:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
378:   PetscFunctionReturn(PETSC_SUCCESS);
379: }

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

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

386:    Not Collective

388:    Input parameters:
389: +  dm - the `DMSWARM`
390: .  celldm - the cell `DM`
391: .  npoints - the number of points to insert in each cell
392: -  xi - the coordinates (defined in the local coordinate system for each cell) to insert

394:  Level: beginner

396:  Notes:
397:  The method will reset any previous defined points within the `DMSWARM`.
398:  Only supported for `DMPLEX`. If you are using a `DMDA` it is recommended to either use
399:  `DMSwarmInsertPointsUsingCellDM()`, or extract and set the coordinates yourself the following code
400: .vb
401:     PetscReal *coor;
402:     DMSwarmGetField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
403:     // user code to define the coordinates here
404:     DMSwarmRestoreField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
405: .ve

407: .seealso: `DMSWARM`, `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
408: @*/
409: PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
410: {
411:   DM        celldm;
412:   PetscBool isDA, isPLEX;

414:   PetscFunctionBegin;
415:   DMSWARMPICVALID(dm);
416:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
417:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
418:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
419:   PetscCheck(!isDA, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMPLEX. Recommended you use DMSwarmInsertPointsUsingCellDM()");
420:   if (isPLEX) {
421:     PetscCall(private_DMSwarmSetPointCoordinatesCellwise_PLEX(dm, celldm, npoints, xi));
422:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
423:   PetscFunctionReturn(PETSC_SUCCESS);
424: }

426: /* Field projection API */
427: extern PetscErrorCode private_DMSwarmProjectFields_DA(DM swarm, DM celldm, PetscInt project_type, PetscInt nfields, DMSwarmDataField dfield[], Vec vecs[]);
428: extern PetscErrorCode private_DMSwarmProjectFields_PLEX(DM swarm, DM celldm, PetscInt project_type, PetscInt nfields, DMSwarmDataField dfield[], Vec vecs[]);

430: /*@C
431:    DMSwarmProjectFields - Project a set of swarm fields onto the cell `DM`

433:    Collective

435:    Input parameters:
436: +  dm - the `DMSWARM`
437: .  nfields - the number of swarm fields to project
438: .  fieldnames - the textual names of the swarm fields to project
439: .  fields - an array of Vec's of length nfields
440: -  reuse - flag indicating whether the array and contents of fields should be re-used or internally allocated

442:    Level: beginner

444:    Notes:
445:    Currently, the only available projection method consists of
446: .vb
447:      phi_i = \sum_{p=0}^{np} N_i(x_p) phi_p dJ / \sum_{p=0}^{np} N_i(x_p) dJ
448:    where phi_p is the swarm field at point p,
449:      N_i() is the cell DM basis function at vertex i,
450:      dJ is the determinant of the cell Jacobian and
451:      phi_i is the projected vertex value of the field phi.
452: .ve

454:    If `reuse` is `PETSC_FALSE`, this function will allocate the array of `Vec`'s, and each individual `Vec`.
455:      The user is responsible for destroying both the array and the individual `Vec` objects.

457:    Only swarm fields registered with data type of `PETSC_REAL` can be projected onto the cell `DM`.

459:    Only swarm fields of block size = 1 can currently be projected.

461:    The only projection methods currently only support the `DMDA` (2D) and `DMPLEX` (triangles 2D).

463: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
464: @*/
465: PETSC_EXTERN PetscErrorCode DMSwarmProjectFields(DM dm, PetscInt nfields, const char *fieldnames[], Vec **fields, PetscBool reuse)
466: {
467:   DM_Swarm         *swarm = (DM_Swarm *)dm->data;
468:   DMSwarmDataField *gfield;
469:   DM                celldm;
470:   PetscBool         isDA, isPLEX;
471:   Vec              *vecs;
472:   PetscInt          f, nvecs;
473:   PetscInt          project_type = 0;

475:   PetscFunctionBegin;
476:   DMSWARMPICVALID(dm);
477:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
478:   PetscCall(PetscMalloc1(nfields, &gfield));
479:   nvecs = 0;
480:   for (f = 0; f < nfields; f++) {
481:     PetscCall(DMSwarmDataBucketGetDMSwarmDataFieldByName(swarm->db, fieldnames[f], &gfield[f]));
482:     PetscCheck(gfield[f]->petsc_type == PETSC_REAL, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Projection only valid for fields using a data type = PETSC_REAL");
483:     PetscCheck(gfield[f]->bs == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Projection only valid for fields with block size = 1");
484:     nvecs += gfield[f]->bs;
485:   }
486:   if (!reuse) {
487:     PetscCall(PetscMalloc1(nvecs, &vecs));
488:     for (f = 0; f < nvecs; f++) {
489:       PetscCall(DMCreateGlobalVector(celldm, &vecs[f]));
490:       PetscCall(PetscObjectSetName((PetscObject)vecs[f], gfield[f]->name));
491:     }
492:   } else {
493:     vecs = *fields;
494:   }

496:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
497:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
498:   if (isDA) {
499:     PetscCall(private_DMSwarmProjectFields_DA(dm, celldm, project_type, nfields, gfield, vecs));
500:   } else if (isPLEX) {
501:     PetscCall(private_DMSwarmProjectFields_PLEX(dm, celldm, project_type, nfields, gfield, vecs));
502:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");

504:   PetscCall(PetscFree(gfield));
505:   if (!reuse) *fields = vecs;
506:   PetscFunctionReturn(PETSC_SUCCESS);
507: }

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

512:    Not Collective

514:    Input parameter:
515: .  dm - the `DMSWARM`

517:    Output parameters:
518: +  ncells - the number of cells in the cell `DM` (optional argument, pass `NULL` to ignore)
519: -  count - array of length ncells containing the number of points per cell

521:    Level: beginner

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

526: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
527: @*/
528: PETSC_EXTERN PetscErrorCode DMSwarmCreatePointPerCellCount(DM dm, PetscInt *ncells, PetscInt **count)
529: {
530:   PetscBool isvalid;
531:   PetscInt  nel;
532:   PetscInt *sum;

534:   PetscFunctionBegin;
535:   PetscCall(DMSwarmSortGetIsValid(dm, &isvalid));
536:   nel = 0;
537:   if (isvalid) {
538:     PetscInt e;

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

542:     PetscCall(PetscMalloc1(nel, &sum));
543:     for (e = 0; e < nel; e++) PetscCall(DMSwarmSortGetNumberOfPointsPerCell(dm, e, &sum[e]));
544:   } else {
545:     DM        celldm;
546:     PetscBool isda, isplex, isshell;
547:     PetscInt  p, npoints;
548:     PetscInt *swarm_cellid;

550:     /* get the number of cells */
551:     PetscCall(DMSwarmGetCellDM(dm, &celldm));
552:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isda));
553:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isplex));
554:     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMSHELL, &isshell));
555:     if (isda) {
556:       PetscInt        _nel, _npe;
557:       const PetscInt *_element;

559:       PetscCall(DMDAGetElements(celldm, &_nel, &_npe, &_element));
560:       nel = _nel;
561:       PetscCall(DMDARestoreElements(celldm, &_nel, &_npe, &_element));
562:     } else if (isplex) {
563:       PetscInt ps, pe;

565:       PetscCall(DMPlexGetHeightStratum(celldm, 0, &ps, &pe));
566:       nel = pe - ps;
567:     } else if (isshell) {
568:       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);

570:       PetscCall(PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
571:       if (method_DMShellGetNumberOfCells) {
572:         PetscCall(method_DMShellGetNumberOfCells(celldm, &nel));
573:       } else
574:         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);");
575:     } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");

577:     PetscCall(PetscMalloc1(nel, &sum));
578:     PetscCall(PetscArrayzero(sum, nel));
579:     PetscCall(DMSwarmGetLocalSize(dm, &npoints));
580:     PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
581:     for (p = 0; p < npoints; p++) {
582:       if (swarm_cellid[p] != DMLOCATEPOINT_POINT_NOT_FOUND) sum[swarm_cellid[p]]++;
583:     }
584:     PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
585:   }
586:   if (ncells) *ncells = nel;
587:   *count = sum;
588:   PetscFunctionReturn(PETSC_SUCCESS);
589: }

591: /*@
592:   DMSwarmGetNumSpecies - Get the number of particle species

594:   Not Collective

596:   Input parameter:
597: . dm - the `DMSWARM`

599:   Output parameters:
600: . Ns - the number of species

602:   Level: intermediate

604: .seealso: `DMSWARM`, `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
605: @*/
606: PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
607: {
608:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

610:   PetscFunctionBegin;
611:   *Ns = swarm->Ns;
612:   PetscFunctionReturn(PETSC_SUCCESS);
613: }

615: /*@
616:   DMSwarmSetNumSpecies - Set the number of particle species

618:   Not Collective

620:   Input parameter:
621: + dm - the `DMSWARM`
622: - Ns - the number of species

624:   Level: intermediate

626: .seealso: `DMSWARM`, `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
627: @*/
628: PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
629: {
630:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

632:   PetscFunctionBegin;
633:   swarm->Ns = Ns;
634:   PetscFunctionReturn(PETSC_SUCCESS);
635: }

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

640:   Not Collective

642:   Input parameter:
643: . dm - the `DMSWARM`

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

648:   Level: intermediate

650: .seealso: `DMSWARM`, `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
651: @*/
652: PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFunc *coordFunc)
653: {
654:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

656:   PetscFunctionBegin;
659:   *coordFunc = swarm->coordFunc;
660:   PetscFunctionReturn(PETSC_SUCCESS);
661: }

663: /*@C
664:   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions

666:   Not Collective

668:   Input parameters:
669: + dm - the `DMSWARM`
670: - coordFunc - the function setting initial particle positions

672:   Level: intermediate

674: .seealso: `DMSWARM`, `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
675: @*/
676: PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFunc coordFunc)
677: {
678:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

680:   PetscFunctionBegin;
683:   swarm->coordFunc = coordFunc;
684:   PetscFunctionReturn(PETSC_SUCCESS);
685: }

687: /*@C
688:   DMSwarmGetCoordinateFunction - Get the function setting initial particle velocities, if it exists

690:   Not Collective

692:   Input parameter:
693: . dm - the `DMSWARM`

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

698:   Level: intermediate

700: .seealso: `DMSWARM`, `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
701: @*/
702: PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFunc *velFunc)
703: {
704:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

706:   PetscFunctionBegin;
709:   *velFunc = swarm->velFunc;
710:   PetscFunctionReturn(PETSC_SUCCESS);
711: }

713: /*@C
714:   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities

716:   Not Collective

718:   Input parameters:
719: + dm - the `DMSWARM`
720: - coordFunc - the function setting initial particle velocities

722:   Level: intermediate

724: .seealso: `DMSWARM`, `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
725: @*/
726: PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFunc velFunc)
727: {
728:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

730:   PetscFunctionBegin;
733:   swarm->velFunc = velFunc;
734:   PetscFunctionReturn(PETSC_SUCCESS);
735: }

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

740:   Not Collective

742:   Input Parameters:
743: + sw      - The `DMSWARM`
744: . N       - The target number of particles
745: - density - The density field for the particle layout, normalized to unity

747:   Level: advanced

749:   Note:
750:   One particle will be created for each species.

752: .seealso: `DMSWARM`, `DMSwarmComputeLocalSizeFromOptions()`
753: @*/
754: PetscErrorCode DMSwarmComputeLocalSize(DM sw, PetscInt N, PetscProbFunc density)
755: {
756:   DM               dm;
757:   PetscQuadrature  quad;
758:   const PetscReal *xq, *wq;
759:   PetscReal       *n_int;
760:   PetscInt        *npc_s, *cellid, Ni;
761:   PetscReal        gmin[3], gmax[3], xi0[3];
762:   PetscInt         Ns, cStart, cEnd, c, dim, d, Nq, q, Np = 0, p, s;
763:   PetscBool        simplex;

765:   PetscFunctionBegin;
766:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
767:   PetscCall(DMSwarmGetCellDM(sw, &dm));
768:   PetscCall(DMGetDimension(dm, &dim));
769:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
770:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
771:   PetscCall(DMPlexIsSimplex(dm, &simplex));
772:   PetscCall(DMGetCoordinatesLocalSetUp(dm));
773:   if (simplex) PetscCall(PetscDTStroudConicalQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
774:   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
775:   PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, &xq, &wq));
776:   PetscCall(PetscCalloc2(Ns, &n_int, (cEnd - cStart) * Ns, &npc_s));
777:   /* Integrate the density function to get the number of particles in each cell */
778:   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
779:   for (c = 0; c < cEnd - cStart; ++c) {
780:     const PetscInt cell = c + cStart;
781:     PetscReal      v0[3], J[9], invJ[9], detJ, detJp = 2. / (gmax[0] - gmin[0]), xr[3], den;

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

791:       for (s = 0; s < Ns; ++s) {
792:         PetscCall(density(xr, NULL, &den));
793:         n_int[s] += (detJp * den) * (detJ * wq[q]) / (PetscReal)Ns;
794:       }
795:     }
796:     for (s = 0; s < Ns; ++s) {
797:       Ni = N;
798:       npc_s[c * Ns + s] += (PetscInt)(Ni * n_int[s]);
799:       Np += npc_s[c * Ns + s];
800:     }
801:   }
802:   PetscCall(PetscQuadratureDestroy(&quad));
803:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
804:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
805:   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
806:     for (s = 0; s < Ns; ++s) {
807:       for (q = 0; q < npc_s[c * Ns + s]; ++q, ++p) cellid[p] = c;
808:     }
809:   }
810:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
811:   PetscCall(PetscFree2(n_int, npc_s));
812:   PetscFunctionReturn(PETSC_SUCCESS);
813: }

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

818:   Not Collective

820:   Input Parameter:
821: . sw - The `DMSWARM`

823:   Level: advanced

825: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`
826: @*/
827: PetscErrorCode DMSwarmComputeLocalSizeFromOptions(DM sw)
828: {
829:   PetscProbFunc pdf;
830:   const char   *prefix;
831:   char          funcname[PETSC_MAX_PATH_LEN];
832:   PetscInt     *N, Ns, dim, n;
833:   PetscBool     flg;
834:   PetscMPIInt   size, rank;

836:   PetscFunctionBegin;
837:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
838:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
839:   PetscCall(PetscCalloc1(size, &N));
840:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
841:   n = size;
842:   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
843:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
844:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
845:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
846:   PetscCall(PetscOptionsString("-dm_swarm_coordinate_function", "Function to determine particle coordinates", "DMSwarmSetCoordinateFunction", funcname, funcname, sizeof(funcname), &flg));
847:   PetscOptionsEnd();
848:   if (flg) {
849:     PetscSimplePointFunc coordFunc;

851:     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
852:     PetscCall(PetscDLSym(NULL, funcname, (void **)&coordFunc));
853:     PetscCheck(coordFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
854:     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
855:     PetscCall(DMSwarmSetLocalSizes(sw, N[rank] * Ns, 0));
856:     PetscCall(DMSwarmSetCoordinateFunction(sw, coordFunc));
857:   } else {
858:     PetscCall(DMGetDimension(sw, &dim));
859:     PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
860:     PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_coordinate_density", &pdf, NULL, NULL));
861:     PetscCall(DMSwarmComputeLocalSize(sw, N[rank], pdf));
862:   }
863:   PetscCall(PetscFree(N));
864:   PetscFunctionReturn(PETSC_SUCCESS);
865: }

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

870:   Not Collective

872:   Input Parameter:
873: . sw - The `DMSWARM`

875:   Level: advanced

877:   Note:
878:   Currently, we randomly place particles in their assigned cell

880: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeVelocities()`
881: @*/
882: PetscErrorCode DMSwarmInitializeCoordinates(DM sw)
883: {
884:   PetscSimplePointFunc coordFunc;
885:   PetscScalar         *weight;
886:   PetscReal           *x;
887:   PetscInt            *species;
888:   void                *ctx;
889:   PetscBool            removePoints = PETSC_TRUE;
890:   PetscDataType        dtype;
891:   PetscInt             Np, p, Ns, dim, d, bs;

893:   PetscFunctionBeginUser;
894:   PetscCall(DMGetDimension(sw, &dim));
895:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
896:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
897:   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));

899:   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, &bs, &dtype, (void **)&x));
900:   PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight));
901:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
902:   if (coordFunc) {
903:     PetscCall(DMGetApplicationContext(sw, &ctx));
904:     for (p = 0; p < Np; ++p) {
905:       PetscScalar X[3];

907:       PetscCall((*coordFunc)(dim, 0., NULL, p, X, ctx));
908:       for (d = 0; d < dim; ++d) x[p * dim + d] = PetscRealPart(X[d]);
909:       weight[p]  = 1.0;
910:       species[p] = p % Ns;
911:     }
912:   } else {
913:     DM          dm;
914:     PetscRandom rnd;
915:     PetscReal   xi0[3];
916:     PetscInt    cStart, cEnd, c;

918:     PetscCall(DMSwarmGetCellDM(sw, &dm));
919:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
920:     PetscCall(DMGetApplicationContext(sw, &ctx));

922:     /* Set particle position randomly in cell, set weights to 1 */
923:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd));
924:     PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
925:     PetscCall(PetscRandomSetFromOptions(rnd));
926:     PetscCall(DMSwarmSortGetAccess(sw));
927:     for (d = 0; d < dim; ++d) xi0[d] = -1.0;
928:     for (c = cStart; c < cEnd; ++c) {
929:       PetscReal v0[3], J[9], invJ[9], detJ;
930:       PetscInt *pidx, Npc, q;

932:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
933:       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ));
934:       for (q = 0; q < Npc; ++q) {
935:         const PetscInt p = pidx[q];
936:         PetscReal      xref[3];

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

941:         weight[p]  = 1.0 / Np;
942:         species[p] = p % Ns;
943:       }
944:       PetscCall(PetscFree(pidx));
945:     }
946:     PetscCall(PetscRandomDestroy(&rnd));
947:     PetscCall(DMSwarmSortRestoreAccess(sw));
948:   }
949:   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
950:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
951:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));

953:   PetscCall(DMSwarmMigrate(sw, removePoints));
954:   PetscCall(DMLocalizeCoordinates(sw));
955:   PetscFunctionReturn(PETSC_SUCCESS);
956: }

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

961:   Collective

963:   Input Parameters:
964: + sw      - The `DMSWARM` object
965: . sampler - A function which uniformly samples the velocity PDF
966: - v0      - The velocity scale for nondimensionalization for each species

968:   Level: advanced

970:   Note:
971:   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.

973: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocitiesFromOptions()`
974: @*/
975: PetscErrorCode DMSwarmInitializeVelocities(DM sw, PetscProbFunc sampler, const PetscReal v0[])
976: {
977:   PetscSimplePointFunc velFunc;
978:   PetscReal           *v;
979:   PetscInt            *species;
980:   void                *ctx;
981:   PetscInt             dim, Np, p;

983:   PetscFunctionBegin;
984:   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));

986:   PetscCall(DMGetDimension(sw, &dim));
987:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
988:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
989:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
990:   if (v0[0] == 0.) {
991:     PetscCall(PetscArrayzero(v, Np * dim));
992:   } else if (velFunc) {
993:     PetscCall(DMGetApplicationContext(sw, &ctx));
994:     for (p = 0; p < Np; ++p) {
995:       PetscInt    s = species[p], d;
996:       PetscScalar vel[3];

998:       PetscCall((*velFunc)(dim, 0., NULL, p, vel, ctx));
999:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
1000:     }
1001:   } else {
1002:     PetscRandom rnd;

1004:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
1005:     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
1006:     PetscCall(PetscRandomSetFromOptions(rnd));

1008:     for (p = 0; p < Np; ++p) {
1009:       PetscInt  s = species[p], d;
1010:       PetscReal a[3], vel[3];

1012:       for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &a[d]));
1013:       PetscCall(sampler(a, NULL, vel));
1014:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * vel[d];
1015:     }
1016:     PetscCall(PetscRandomDestroy(&rnd));
1017:   }
1018:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1019:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1020:   PetscFunctionReturn(PETSC_SUCCESS);
1021: }

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

1026:   Collective

1028:   Input Parameters:
1029: + sw      - The `DMSWARM` object
1030: - v0      - The velocity scale for nondimensionalization for each species

1032:   Level: advanced

1034: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
1035: @*/
1036: PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
1037: {
1038:   PetscProbFunc sampler;
1039:   PetscInt      dim;
1040:   const char   *prefix;
1041:   char          funcname[PETSC_MAX_PATH_LEN];
1042:   PetscBool     flg;

1044:   PetscFunctionBegin;
1045:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
1046:   PetscCall(PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg));
1047:   PetscOptionsEnd();
1048:   if (flg) {
1049:     PetscSimplePointFunc velFunc;

1051:     PetscCall(PetscDLSym(NULL, funcname, (void **)&velFunc));
1052:     PetscCheck(velFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
1053:     PetscCall(DMSwarmSetVelocityFunction(sw, velFunc));
1054:   }
1055:   PetscCall(DMGetDimension(sw, &dim));
1056:   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
1057:   PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_velocity_density", NULL, NULL, &sampler));
1058:   PetscCall(DMSwarmInitializeVelocities(sw, sampler, v0));
1059:   PetscFunctionReturn(PETSC_SUCCESS);
1060: }