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; \
 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 on dm

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

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

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

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

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

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

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

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

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

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

174:   PetscSFDestroy(&sfcell);
175:   VecDestroy(&pos);
176:   return 0;
177: }

179: /*@C
180:    DMSwarmSetPointCoordinates - Set point coordinates in a DMSwarm from a user defined list

182:    Collective on dm

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

191:    Level: beginner

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

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

219:   DMSWARMPICVALID(dm);
220:   PetscObjectGetComm((PetscObject)dm, &comm);
221:   MPI_Comm_rank(comm, &rank);

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

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

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

253:   /* determine the number of points living in the bounding box */
254:   n_estimate = 0;
255:   for (i = 0; i < my_npoints; i++) {
256:     PetscBool point_inside = PETSC_TRUE;

258:     for (b = 0; b < bs; b++) {
259:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
260:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
261:     }
262:     if (point_inside) n_estimate++;
263:   }

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

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

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

287:   /* locate points */
288:   DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell);

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

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

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

324:   if (redundant) {
325:     if (rank > 0) PetscFree(my_coor);
326:   }
327:   PetscSFDestroy(&sfcell);
328:   VecDestroy(&pos);
329:   return 0;
330: }

332: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
333: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);

335: /*@C
336:    DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell

338:    Not collective

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

345:    Level: beginner

347:    Notes:

349:    The insert method will reset any previous defined points within the DMSwarm.

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

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

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

365:   DMSWARMPICVALID(dm);
366:   DMSwarmGetCellDM(dm, &celldm);
367:   PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA);
368:   PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX);
369:   if (isDA) {
370:     private_DMSwarmInsertPointsUsingCellDM_DA(dm, celldm, layout_type, fill_param);
371:   } else if (isPLEX) {
372:     private_DMSwarmInsertPointsUsingCellDM_PLEX(dm, celldm, layout_type, fill_param);
373:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
374:   return 0;
375: }

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

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

382:    Not collective

384:    Input parameters:
385: +  dm - the DMSwarm
386: .  celldm - the cell DM
387: .  npoints - the number of points to insert in each cell
388: -  xi - the coordinates (defined in the local coordinate system for each cell) to insert

390:  Level: beginner

392:  Notes:
393:  The method will reset any previous defined points within the DMSwarm.
394:  Only supported for DMPLEX. If you are using a DMDA it is recommended to either use
395:  DMSwarmInsertPointsUsingCellDM(), or extract and set the coordinates yourself the following code

397: $    PetscReal *coor;
398: $    DMSwarmGetField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
399: $    // user code to define the coordinates here
400: $    DMSwarmRestoreField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);

402: .seealso: `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
403: @*/
404: PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
405: {
406:   DM        celldm;
407:   PetscBool isDA, isPLEX;

409:   DMSWARMPICVALID(dm);
410:   DMSwarmGetCellDM(dm, &celldm);
411:   PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA);
412:   PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX);
414:   if (isPLEX) {
415:     private_DMSwarmSetPointCoordinatesCellwise_PLEX(dm, celldm, npoints, xi);
416:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
417:   return 0;
418: }

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

424: /*@C
425:    DMSwarmProjectFields - Project a set of swarm fields onto the cell DM

427:    Collective on dm

429:    Input parameters:
430: +  dm - the DMSwarm
431: .  nfields - the number of swarm fields to project
432: .  fieldnames - the textual names of the swarm fields to project
433: .  fields - an array of Vec's of length nfields
434: -  reuse - flag indicating whether the array and contents of fields should be re-used or internally allocated

436:    Currently, the only available projection method consists of
437:      phi_i = \sum_{p=0}^{np} N_i(x_p) phi_p dJ / \sum_{p=0}^{np} N_i(x_p) dJ
438:    where phi_p is the swarm field at point p,
439:      N_i() is the cell DM basis function at vertex i,
440:      dJ is the determinant of the cell Jacobian and
441:      phi_i is the projected vertex value of the field phi.

443:    Level: beginner

445:    Notes:

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

450:    Only swarm fields registered with data type = PETSC_REAL can be projected onto the cell DM.

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

454:    The only projection methods currently only support the DA (2D) and PLEX (triangles 2D).

456: .seealso: `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
457: @*/
458: PETSC_EXTERN PetscErrorCode DMSwarmProjectFields(DM dm, PetscInt nfields, const char *fieldnames[], Vec **fields, PetscBool reuse)
459: {
460:   DM_Swarm         *swarm = (DM_Swarm *)dm->data;
461:   DMSwarmDataField *gfield;
462:   DM                celldm;
463:   PetscBool         isDA, isPLEX;
464:   Vec              *vecs;
465:   PetscInt          f, nvecs;
466:   PetscInt          project_type = 0;

468:   DMSWARMPICVALID(dm);
469:   DMSwarmGetCellDM(dm, &celldm);
470:   PetscMalloc1(nfields, &gfield);
471:   nvecs = 0;
472:   for (f = 0; f < nfields; f++) {
473:     DMSwarmDataBucketGetDMSwarmDataFieldByName(swarm->db, fieldnames[f], &gfield[f]);
476:     nvecs += gfield[f]->bs;
477:   }
478:   if (!reuse) {
479:     PetscMalloc1(nvecs, &vecs);
480:     for (f = 0; f < nvecs; f++) {
481:       DMCreateGlobalVector(celldm, &vecs[f]);
482:       PetscObjectSetName((PetscObject)vecs[f], gfield[f]->name);
483:     }
484:   } else {
485:     vecs = *fields;
486:   }

488:   PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA);
489:   PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX);
490:   if (isDA) {
491:     private_DMSwarmProjectFields_DA(dm, celldm, project_type, nfields, gfield, vecs);
492:   } else if (isPLEX) {
493:     private_DMSwarmProjectFields_PLEX(dm, celldm, project_type, nfields, gfield, vecs);
494:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");

496:   PetscFree(gfield);
497:   if (!reuse) *fields = vecs;
498:   return 0;
499: }

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

504:    Not collective

506:    Input parameter:
507: .  dm - the DMSwarm

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

513:    Level: beginner

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

518: .seealso: `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
519: @*/
520: PETSC_EXTERN PetscErrorCode DMSwarmCreatePointPerCellCount(DM dm, PetscInt *ncells, PetscInt **count)
521: {
522:   PetscBool isvalid;
523:   PetscInt  nel;
524:   PetscInt *sum;

526:   DMSwarmSortGetIsValid(dm, &isvalid);
527:   nel = 0;
528:   if (isvalid) {
529:     PetscInt e;

531:     DMSwarmSortGetSizes(dm, &nel, NULL);

533:     PetscMalloc1(nel, &sum);
534:     for (e = 0; e < nel; e++) DMSwarmSortGetNumberOfPointsPerCell(dm, e, &sum[e]);
535:   } else {
536:     DM        celldm;
537:     PetscBool isda, isplex, isshell;
538:     PetscInt  p, npoints;
539:     PetscInt *swarm_cellid;

541:     /* get the number of cells */
542:     DMSwarmGetCellDM(dm, &celldm);
543:     PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isda);
544:     PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isplex);
545:     PetscObjectTypeCompare((PetscObject)celldm, DMSHELL, &isshell);
546:     if (isda) {
547:       PetscInt        _nel, _npe;
548:       const PetscInt *_element;

550:       DMDAGetElements(celldm, &_nel, &_npe, &_element);
551:       nel = _nel;
552:       DMDARestoreElements(celldm, &_nel, &_npe, &_element);
553:     } else if (isplex) {
554:       PetscInt ps, pe;

556:       DMPlexGetHeightStratum(celldm, 0, &ps, &pe);
557:       nel = pe - ps;
558:     } else if (isshell) {
559:       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);

561:       PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells);
562:       if (method_DMShellGetNumberOfCells) {
563:         method_DMShellGetNumberOfCells(celldm, &nel);
564:       } else
565:         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);");
566:     } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");

568:     PetscMalloc1(nel, &sum);
569:     PetscArrayzero(sum, nel);
570:     DMSwarmGetLocalSize(dm, &npoints);
571:     DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid);
572:     for (p = 0; p < npoints; p++) {
573:       if (swarm_cellid[p] != DMLOCATEPOINT_POINT_NOT_FOUND) sum[swarm_cellid[p]]++;
574:     }
575:     DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid);
576:   }
577:   if (ncells) *ncells = nel;
578:   *count = sum;
579:   return 0;
580: }

582: /*@
583:   DMSwarmGetNumSpecies - Get the number of particle species

585:   Not collective

587:   Input parameter:
588: . dm - the DMSwarm

590:   Output parameters:
591: . Ns - the number of species

593:   Level: intermediate

595: .seealso: `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
596: @*/
597: PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
598: {
599:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

601:   *Ns = swarm->Ns;
602:   return 0;
603: }

605: /*@
606:   DMSwarmSetNumSpecies - Set the number of particle species

608:   Not collective

610:   Input parameter:
611: + dm - the DMSwarm
612: - Ns - the number of species

614:   Level: intermediate

616: .seealso: `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
617: @*/
618: PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
619: {
620:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

622:   swarm->Ns = Ns;
623:   return 0;
624: }

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

629:   Not collective

631:   Input parameter:
632: . dm - the DMSwarm

634:   Output Parameter:
635: . coordFunc - the function setting initial particle positions, or NULL

637:   Level: intermediate

639: .seealso: `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
640: @*/
641: PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFunc *coordFunc)
642: {
643:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

647:   *coordFunc = swarm->coordFunc;
648:   return 0;
649: }

651: /*@C
652:   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions

654:   Not collective

656:   Input parameters:
657: + dm - the DMSwarm
658: - coordFunc - the function setting initial particle positions

660:   Level: intermediate

662: .seealso: `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
663: @*/
664: PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFunc coordFunc)
665: {
666:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

670:   swarm->coordFunc = coordFunc;
671:   return 0;
672: }

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

677:   Not collective

679:   Input parameter:
680: . dm - the DMSwarm

682:   Output Parameter:
683: . velFunc - the function setting initial particle velocities, or NULL

685:   Level: intermediate

687: .seealso: `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
688: @*/
689: PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFunc *velFunc)
690: {
691:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

695:   *velFunc = swarm->velFunc;
696:   return 0;
697: }

699: /*@C
700:   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities

702:   Not collective

704:   Input parameters:
705: + dm - the DMSwarm
706: - coordFunc - the function setting initial particle velocities

708:   Level: intermediate

710: .seealso: `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
711: @*/
712: PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFunc velFunc)
713: {
714:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

718:   swarm->velFunc = velFunc;
719:   return 0;
720: }

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

725:   Not collective

727:   Input Parameters:
728: + sw      - The DMSwarm
729: . N       - The target number of particles
730: - density - The density field for the particle layout, normalized to unity

732:   Note: One particle will be created for each species.

734:   Level: advanced

736: .seealso: `DMSwarmComputeLocalSizeFromOptions()`
737: @*/
738: PetscErrorCode DMSwarmComputeLocalSize(DM sw, PetscInt N, PetscProbFunc density)
739: {
740:   DM               dm;
741:   PetscQuadrature  quad;
742:   const PetscReal *xq, *wq;
743:   PetscInt        *npc, *cellid;
744:   PetscReal        xi0[3];
745:   PetscInt         Ns, cStart, cEnd, c, dim, d, Nq, q, Np = 0, p;
746:   PetscBool        simplex;

748:   DMSwarmGetNumSpecies(sw, &Ns);
749:   DMSwarmGetCellDM(sw, &dm);
750:   DMGetDimension(dm, &dim);
751:   DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd);
752:   DMPlexIsSimplex(dm, &simplex);
753:   DMGetCoordinatesLocalSetUp(dm);
754:   if (simplex) PetscDTStroudConicalQuadrature(dim, 1, 5, -1.0, 1.0, &quad);
755:   else PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad);
756:   PetscQuadratureGetData(quad, NULL, NULL, &Nq, &xq, &wq);
757:   PetscMalloc1(cEnd - cStart, &npc);
758:   /* Integrate the density function to get the number of particles in each cell */
759:   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
760:   for (c = 0; c < cEnd - cStart; ++c) {
761:     const PetscInt cell = c + cStart;
762:     PetscReal      v0[3], J[9], invJ[9], detJ;
763:     PetscReal      n_int = 0.;

765:     DMPlexComputeCellGeometryFEM(dm, cell, NULL, v0, J, invJ, &detJ);
766:     for (q = 0; q < Nq; ++q) {
767:       PetscReal xr[3], den[3];

769:       CoordinatesRefToReal(dim, dim, xi0, v0, J, &xq[q * dim], xr);
770:       density(xr, NULL, den);
771:       n_int += den[0] * wq[q];
772:     }
773:     npc[c] = (PetscInt)(N * n_int);
774:     npc[c] *= Ns;
775:     Np += npc[c];
776:   }
777:   PetscQuadratureDestroy(&quad);
778:   DMSwarmSetLocalSizes(sw, Np, 0);

780:   DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid);
781:   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
782:     for (q = 0; q < npc[c]; ++q, ++p) cellid[p] = c;
783:   }
784:   DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid);
785:   PetscFree(npc);
786:   return 0;
787: }

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

792:   Not collective

794:   Input Parameters:
795: , sw - The DMSwarm

797:   Level: advanced

799: .seealso: `DMSwarmComputeLocalSize()`
800: @*/
801: PetscErrorCode DMSwarmComputeLocalSizeFromOptions(DM sw)
802: {
803:   PetscProbFunc pdf;
804:   const char   *prefix;
805:   char          funcname[PETSC_MAX_PATH_LEN];
806:   PetscInt     *N, Ns, dim, n;
807:   PetscBool     flg;
808:   PetscMPIInt   size, rank;

810:   MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size);
811:   MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank);
812:   PetscCalloc1(size, &N);
813:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
814:   n = size;
815:   PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL);
816:   DMSwarmGetNumSpecies(sw, &Ns);
817:   PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg);
818:   if (flg) DMSwarmSetNumSpecies(sw, Ns);
819:   PetscOptionsString("-dm_swarm_coordinate_function", "Function to determine particle coordinates", "DMSwarmSetCoordinateFunction", funcname, funcname, sizeof(funcname), &flg);
820:   PetscOptionsEnd();
821:   if (flg) {
822:     PetscSimplePointFunc coordFunc;

824:     DMSwarmGetNumSpecies(sw, &Ns);
825:     PetscDLSym(NULL, funcname, (void **)&coordFunc);
827:     DMSwarmGetNumSpecies(sw, &Ns);
828:     DMSwarmSetLocalSizes(sw, N[rank] * Ns, 0);
829:     DMSwarmSetCoordinateFunction(sw, coordFunc);
830:   } else {
831:     DMGetDimension(sw, &dim);
832:     PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix);
833:     PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_coordinate_density", &pdf, NULL, NULL);
834:     DMSwarmComputeLocalSize(sw, N[rank], pdf);
835:   }
836:   PetscFree(N);
837:   return 0;
838: }

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

843:   Not collective

845:   Input Parameters:
846: , sw - The DMSwarm

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

850:   Level: advanced

852: .seealso: `DMSwarmComputeLocalSize()`, `DMSwarmInitializeVelocities()`
853: @*/
854: PetscErrorCode DMSwarmInitializeCoordinates(DM sw)
855: {
856:   PetscSimplePointFunc coordFunc;
857:   PetscScalar         *weight;
858:   PetscReal           *x;
859:   PetscInt            *species;
860:   void                *ctx;
861:   PetscBool            removePoints = PETSC_TRUE;
862:   PetscDataType        dtype;
863:   PetscInt             Np, p, Ns, dim, d, bs;

866:   DMGetDimension(sw, &dim);
867:   DMSwarmGetLocalSize(sw, &Np);
868:   DMSwarmGetNumSpecies(sw, &Ns);
869:   DMSwarmGetCoordinateFunction(sw, &coordFunc);

871:   DMSwarmGetField(sw, DMSwarmPICField_coor, &bs, &dtype, (void **)&x);
872:   DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight);
873:   DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species);
874:   if (coordFunc) {
875:     DMGetApplicationContext(sw, &ctx);
876:     for (p = 0; p < Np; ++p) {
877:       PetscScalar X[3];

879:       (*coordFunc)(dim, 0., NULL, p, X, ctx);
880:       for (d = 0; d < dim; ++d) x[p * dim + d] = PetscRealPart(X[d]);
881:       weight[p]  = 1.0;
882:       species[p] = p % Ns;
883:     }
884:   } else {
885:     DM          dm;
886:     PetscRandom rnd;
887:     PetscReal   xi0[3];
888:     PetscInt    cStart, cEnd, c;

890:     DMSwarmGetCellDM(sw, &dm);
891:     DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd);

893:     /* Set particle position randomly in cell, set weights to 1 */
894:     PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd);
895:     PetscRandomSetInterval(rnd, -1.0, 1.0);
896:     PetscRandomSetFromOptions(rnd);
897:     DMSwarmSortGetAccess(sw);
898:     for (d = 0; d < dim; ++d) xi0[d] = -1.0;
899:     for (c = cStart; c < cEnd; ++c) {
900:       PetscReal v0[3], J[9], invJ[9], detJ;
901:       PetscInt *pidx, Npc, q;

903:       DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx);
904:       DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ);
905:       for (q = 0; q < Npc; ++q) {
906:         const PetscInt p = pidx[q];
907:         PetscReal      xref[3];

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

912:         weight[p]  = 1.0;
913:         species[p] = p % Ns;
914:       }
915:       PetscFree(pidx);
916:     }
917:     PetscRandomDestroy(&rnd);
918:     DMSwarmSortRestoreAccess(sw);
919:   }
920:   DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x);
921:   DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight);
922:   DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species);

924:   DMSwarmMigrate(sw, removePoints);
925:   DMLocalizeCoordinates(sw);
926:   return 0;
927: }

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

932:   Collective on dm

934:   Input Parameters:
935: + sw      - The DMSwarm object
936: . sampler - A function which uniformly samples the velocity PDF
937: - v0      - The velocity scale for nondimensionalization for each species

939:   Note: 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.

941:   Level: advanced

943: .seealso: `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocitiesFromOptions()`
944: @*/
945: PetscErrorCode DMSwarmInitializeVelocities(DM sw, PetscProbFunc sampler, const PetscReal v0[])
946: {
947:   PetscSimplePointFunc velFunc;
948:   PetscReal           *v;
949:   PetscInt            *species;
950:   void                *ctx;
951:   PetscInt             dim, Np, p;

953:   DMSwarmGetVelocityFunction(sw, &velFunc);

955:   DMGetDimension(sw, &dim);
956:   DMSwarmGetLocalSize(sw, &Np);
957:   DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v);
958:   DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species);
959:   if (v0[0] == 0.) {
960:     PetscArrayzero(v, Np * dim);
961:   } else if (velFunc) {
962:     DMGetApplicationContext(sw, &ctx);
963:     for (p = 0; p < Np; ++p) {
964:       PetscInt    s = species[p], d;
965:       PetscScalar vel[3];

967:       (*velFunc)(dim, 0., NULL, p, vel, ctx);
968:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
969:     }
970:   } else {
971:     PetscRandom rnd;

973:     PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd);
974:     PetscRandomSetInterval(rnd, 0, 1.);
975:     PetscRandomSetFromOptions(rnd);

977:     for (p = 0; p < Np; ++p) {
978:       PetscInt  s = species[p], d;
979:       PetscReal a[3], vel[3];

981:       for (d = 0; d < dim; ++d) PetscRandomGetValueReal(rnd, &a[d]);
982:       sampler(a, NULL, vel);
983:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * vel[d];
984:     }
985:     PetscRandomDestroy(&rnd);
986:   }
987:   DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v);
988:   DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species);
989:   return 0;
990: }

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

995:   Collective on dm

997:   Input Parameters:
998: + sw      - The DMSwarm object
999: - v0      - The velocity scale for nondimensionalization for each species

1001:   Level: advanced

1003: .seealso: `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
1004: @*/
1005: PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
1006: {
1007:   PetscProbFunc sampler;
1008:   PetscInt      dim;
1009:   const char   *prefix;
1010:   char          funcname[PETSC_MAX_PATH_LEN];
1011:   PetscBool     flg;

1013:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
1014:   PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg);
1015:   PetscOptionsEnd();
1016:   if (flg) {
1017:     PetscSimplePointFunc velFunc;

1019:     PetscDLSym(NULL, funcname, (void **)&velFunc);
1021:     DMSwarmSetVelocityFunction(sw, velFunc);
1022:   }
1023:   DMGetDimension(sw, &dim);
1024:   PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix);
1025:   PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_velocity_density", NULL, NULL, &sampler);
1026:   DMSwarmInitializeVelocities(sw, sampler, v0);
1027:   return 0;
1028: }