Actual source code: swarmpic.c

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

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

 10: PetscClassId DMSWARMCELLDM_CLASSID;

 12: /*@
 13:   DMSwarmCellDMDestroy - destroy a `DMSwarmCellDM`

 15:   Collective

 17:   Input Parameter:
 18: . celldm - address of `DMSwarmCellDM`

 20:   Level: advanced

 22: .seealso: `DMSwarmCellDM`, `DMSwarmCellDMCreate()`
 23: @*/
 24: PetscErrorCode DMSwarmCellDMDestroy(DMSwarmCellDM *celldm)
 25: {
 26:   PetscFunctionBegin;
 27:   if (!*celldm) PetscFunctionReturn(PETSC_SUCCESS);
 29:   if (--((PetscObject)*celldm)->refct > 0) {
 30:     *celldm = NULL;
 31:     PetscFunctionReturn(PETSC_SUCCESS);
 32:   }
 33:   PetscTryTypeMethod(*celldm, destroy);
 34:   for (PetscInt f = 0; f < (*celldm)->Nf; ++f) PetscCall(PetscFree((*celldm)->dmFields[f]));
 35:   PetscCall(PetscFree((*celldm)->dmFields));
 36:   for (PetscInt f = 0; f < (*celldm)->Nfc; ++f) PetscCall(PetscFree((*celldm)->coordFields[f]));
 37:   PetscCall(PetscFree((*celldm)->coordFields));
 38:   PetscCall(PetscFree((*celldm)->cellid));
 39:   PetscCall(DMSwarmSortDestroy(&(*celldm)->sort));
 40:   PetscCall(DMDestroy(&(*celldm)->dm));
 41:   PetscCall(PetscHeaderDestroy(celldm));
 42:   PetscFunctionReturn(PETSC_SUCCESS);
 43: }

 45: /*@
 46:   DMSwarmCellDMView - view a `DMSwarmCellDM`

 48:   Collective

 50:   Input Parameters:
 51: + celldm - `DMSwarmCellDM`
 52: - viewer - viewer to display field, for example `PETSC_VIEWER_STDOUT_WORLD`

 54:   Level: advanced

 56: .seealso: `DMSwarmCellDM`, `DMSwarmCellDMCreate()`
 57: @*/
 58: PetscErrorCode DMSwarmCellDMView(DMSwarmCellDM celldm, PetscViewer viewer)
 59: {
 60:   PetscBool isascii;

 62:   PetscFunctionBegin;
 64:   if (!viewer) PetscCall(PetscViewerASCIIGetStdout(PetscObjectComm((PetscObject)celldm), &viewer));
 66:   PetscCheckSameComm(celldm, 1, viewer, 2);
 67:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
 68:   if (isascii) {
 69:     PetscCall(PetscObjectPrintClassNamePrefixType((PetscObject)celldm, viewer));
 70:     PetscCall(PetscViewerASCIIPushTab(viewer));
 71:     PetscCall(PetscViewerASCIIPrintf(viewer, "solution field%s:", celldm->Nf > 1 ? "s" : ""));
 72:     for (PetscInt f = 0; f < celldm->Nf; ++f) PetscCall(PetscViewerASCIIPrintf(viewer, " %s", celldm->dmFields[f]));
 73:     PetscCall(PetscViewerASCIIPrintf(viewer, "\n"));
 74:     PetscCall(PetscViewerASCIIPrintf(viewer, "coordinate field%s:", celldm->Nfc > 1 ? "s" : ""));
 75:     for (PetscInt f = 0; f < celldm->Nfc; ++f) PetscCall(PetscViewerASCIIPrintf(viewer, " %s", celldm->coordFields[f]));
 76:     PetscCall(PetscViewerASCIIPrintf(viewer, "\n"));
 77:     PetscCall(PetscViewerPushFormat(viewer, PETSC_VIEWER_DEFAULT));
 78:     PetscCall(DMView(celldm->dm, viewer));
 79:     PetscCall(PetscViewerPopFormat(viewer));
 80:   }
 81:   PetscTryTypeMethod(celldm, view, viewer);
 82:   if (isascii) PetscCall(PetscViewerASCIIPopTab(viewer));
 83:   PetscFunctionReturn(PETSC_SUCCESS);
 84: }

 86: /*@
 87:   DMSwarmCellDMGetDM - Returns the background `DM` for the `DMSwarm`

 89:   Not Collective

 91:   Input Parameter:
 92: . celldm - The `DMSwarmCellDM` object

 94:   Output Parameter:
 95: . dm - The `DM` object

 97:   Level: intermediate

 99: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`
100: @*/
101: PetscErrorCode DMSwarmCellDMGetDM(DMSwarmCellDM celldm, DM *dm)
102: {
103:   PetscFunctionBegin;
105:   PetscAssertPointer(dm, 2);
106:   *dm = celldm->dm;
107:   PetscFunctionReturn(PETSC_SUCCESS);
108: }

110: /*@C
111:   DMSwarmCellDMGetFields - Returns the `DM` fields for the `DMSwarm`

113:   Not Collective

115:   Input Parameter:
116: . celldm - The `DMSwarmCellDM` object

118:   Output Parameters:
119: + Nf    - The number of fields
120: - names - The array of field names in the `DMSWARM`

122:   Level: intermediate

124: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`
125: @*/
126: PetscErrorCode DMSwarmCellDMGetFields(DMSwarmCellDM celldm, PetscInt *Nf, const char **names[])
127: {
128:   PetscFunctionBegin;
130:   if (Nf) {
131:     PetscAssertPointer(Nf, 2);
132:     *Nf = celldm->Nf;
133:   }
134:   if (names) {
135:     PetscAssertPointer(names, 3);
136:     *names = (const char **)celldm->dmFields;
137:   }
138:   PetscFunctionReturn(PETSC_SUCCESS);
139: }

141: /*@C
142:   DMSwarmCellDMGetCoordinateFields - Returns the `DM` coordinate fields for the `DMSwarm`

144:   Not Collective

146:   Input Parameter:
147: . celldm - The `DMSwarmCellDM` object

149:   Output Parameters:
150: + Nfc   - The number of coordinate fields
151: - names - The array of coordinate field names in the `DMSWARM`

153:   Level: intermediate

155: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`
156: @*/
157: PetscErrorCode DMSwarmCellDMGetCoordinateFields(DMSwarmCellDM celldm, PetscInt *Nfc, const char **names[])
158: {
159:   PetscFunctionBegin;
161:   if (Nfc) {
162:     PetscAssertPointer(Nfc, 2);
163:     *Nfc = celldm->Nfc;
164:   }
165:   if (names) {
166:     PetscAssertPointer(names, 3);
167:     *names = (const char **)celldm->coordFields;
168:   }
169:   PetscFunctionReturn(PETSC_SUCCESS);
170: }

172: /*@C
173:   DMSwarmCellDMGetCellID - Returns the cell id field name for the `DMSwarm`

175:   Not Collective

177:   Input Parameter:
178: . celldm - The `DMSwarmCellDM` object

180:   Output Parameters:
181: . cellid - The cell id field name in the `DMSWARM`

183:   Level: intermediate

185: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`
186: @*/
187: PetscErrorCode DMSwarmCellDMGetCellID(DMSwarmCellDM celldm, const char *cellid[])
188: {
189:   PetscFunctionBegin;
191:   PetscAssertPointer(cellid, 2);
192:   *cellid = celldm->cellid;
193:   PetscFunctionReturn(PETSC_SUCCESS);
194: }

196: /*@C
197:   DMSwarmCellDMGetSort - Returns the sort context over the active `DMSwarmCellDM` for the `DMSwarm`

199:   Not Collective

201:   Input Parameter:
202: . celldm - The `DMSwarmCellDM` object

204:   Output Parameter:
205: . sort - The `DMSwarmSort` object

207:   Level: intermediate

209: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`, `DMSwarmSortGetAccess()`, `DMSwarmSortDestroy()`
210: @*/
211: PetscErrorCode DMSwarmCellDMGetSort(DMSwarmCellDM celldm, DMSwarmSort *sort)
212: {
213:   PetscFunctionBegin;
215:   PetscAssertPointer(sort, 2);
216:   *sort = celldm->sort;
217:   PetscFunctionReturn(PETSC_SUCCESS);
218: }

220: /*@C
221:   DMSwarmCellDMSetSort - Sets the sort context over the active `DMSwarmCellDM` for the `DMSwarm`

223:   Not Collective

225:   Input Parameters:
226: + celldm - The `DMSwarmCellDM` object
227: - sort   - The `DMSwarmSort` object, or `NULL` to clear the context

229:   Level: intermediate

231:   Note:
232:   User should destroy `sort` afterward with `DMSwarmSortDestroy()`, as the `DMSwarmCellDM` will hold a reference.

234: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`, `DMSwarmSortGetAccess()`, `DMSwarmSortDestroy()`
235: @*/
236: PetscErrorCode DMSwarmCellDMSetSort(DMSwarmCellDM celldm, DMSwarmSort sort)
237: {
238:   PetscFunctionBegin;
240:   if (sort) PetscAssertPointer(sort, 2);
241:   PetscCall(PetscObjectReference((PetscObject)sort));
242:   PetscCall(DMSwarmSortDestroy(&celldm->sort));
243:   celldm->sort = sort;
244:   PetscFunctionReturn(PETSC_SUCCESS);
245: }

247: /*@
248:   DMSwarmCellDMGetBlockSize - Returns the total blocksize for the `DM` fields

250:   Not Collective

252:   Input Parameters:
253: + celldm - The `DMSwarmCellDM` object
254: - sw     - The `DMSwarm` object

256:   Output Parameter:
257: . bs - The total block size

259:   Level: intermediate

261: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`
262: @*/
263: PetscErrorCode DMSwarmCellDMGetBlockSize(DMSwarmCellDM celldm, DM sw, PetscInt *bs)
264: {
265:   PetscFunctionBegin;
268:   PetscAssertPointer(bs, 3);
269:   *bs = 0;
270:   for (PetscInt f = 0; f < celldm->Nf; ++f) {
271:     PetscInt fbs;

273:     PetscCall(DMSwarmGetFieldInfo(sw, celldm->dmFields[f], &fbs, NULL));
274:     *bs += fbs;
275:   }
276:   PetscFunctionReturn(PETSC_SUCCESS);
277: }

279: /*@
280:   DMSwarmCellDMCreate - create a `DMSwarmCellDM`

282:   Collective

284:   Input Parameters:
285: + dm          - The background `DM` for the `DMSwarm`
286: . Nf          - The number of swarm fields defined over `dm`
287: . dmFields    - The swarm field names for the `dm` fields
288: . Nfc         - The number of swarm fields to use for coordinates over `dm`
289: - coordFields - The swarm field names for the `dm` coordinate fields

291:   Output Parameter:
292: . celldm - The new `DMSwarmCellDM`

294:   Level: advanced

296: .seealso: `DMSwarmCellDM`, `DMSWARM`, `DMSetType()`
297: @*/
298: PetscErrorCode DMSwarmCellDMCreate(DM dm, PetscInt Nf, const char *dmFields[], PetscInt Nfc, const char *coordFields[], DMSwarmCellDM *celldm)
299: {
300:   DMSwarmCellDM b;
301:   const char   *name;
302:   char          cellid[PETSC_MAX_PATH_LEN];

304:   PetscFunctionBegin;
306:   if (Nf) PetscAssertPointer(dmFields, 3);
307:   if (Nfc) PetscAssertPointer(coordFields, 5);
308:   PetscCall(DMInitializePackage());

310:   PetscCall(PetscHeaderCreate(b, DMSWARMCELLDM_CLASSID, "DMSwarmCellDM", "Background DM for a Swarm", "DM", PetscObjectComm((PetscObject)dm), DMSwarmCellDMDestroy, DMSwarmCellDMView));
311:   PetscCall(PetscObjectGetName((PetscObject)dm, &name));
312:   PetscCall(PetscObjectSetName((PetscObject)b, name));
313:   PetscCall(PetscObjectReference((PetscObject)dm));
314:   b->dm  = dm;
315:   b->Nf  = Nf;
316:   b->Nfc = Nfc;
317:   PetscCall(PetscMalloc1(b->Nf, &b->dmFields));
318:   for (PetscInt f = 0; f < b->Nf; ++f) PetscCall(PetscStrallocpy(dmFields[f], &b->dmFields[f]));
319:   PetscCall(PetscMalloc1(b->Nfc, &b->coordFields));
320:   for (PetscInt f = 0; f < b->Nfc; ++f) PetscCall(PetscStrallocpy(coordFields[f], &b->coordFields[f]));
321:   PetscCall(PetscSNPrintf(cellid, PETSC_MAX_PATH_LEN, "%s_cellid", name));
322:   PetscCall(PetscStrallocpy(cellid, &b->cellid));
323:   *celldm = b;
324:   PetscFunctionReturn(PETSC_SUCCESS);
325: }

327: /* Coordinate insertition/addition API */
328: /*@
329:   DMSwarmSetPointsUniformCoordinates - Set point coordinates in a `DMSWARM` on a regular (ijk) grid

331:   Collective

333:   Input Parameters:
334: + sw      - the `DMSWARM`
335: . min     - minimum coordinate values in the x, y, z directions (array of length dim)
336: . max     - maximum coordinate values in the x, y, z directions (array of length dim)
337: . npoints - number of points in each spatial direction (array of length dim)
338: - mode    - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)

340:   Level: beginner

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

347: .seealso: `DM`, `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
348: @*/
349: PetscErrorCode DMSwarmSetPointsUniformCoordinates(DM sw, PetscReal min[], PetscReal max[], PetscInt npoints[], InsertMode mode)
350: {
351:   PetscReal          lmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
352:   PetscReal          lmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
353:   PetscInt           i, j, k, bs, b, n_estimate, n_curr, n_new_est, p, n_found, Nfc;
354:   const PetscScalar *_coor;
355:   DMSwarmCellDM      celldm;
356:   DM                 dm;
357:   PetscReal          dx[3];
358:   PetscInt           _npoints[] = {0, 0, 1};
359:   Vec                pos;
360:   PetscScalar       *_pos;
361:   PetscReal         *swarm_coor;
362:   PetscInt          *swarm_cellid;
363:   PetscSF            sfcell = NULL;
364:   const PetscSFNode *LA_sfcell;
365:   const char       **coordFields, *cellid;

367:   PetscFunctionBegin;
368:   DMSWARMPICVALID(sw);
369:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
370:   PetscCall(DMSwarmCellDMGetCoordinateFields(celldm, &Nfc, &coordFields));
371:   PetscCheck(Nfc == 1, PetscObjectComm((PetscObject)sw), PETSC_ERR_SUP, "We only support a single coordinate field right now, not %" PetscInt_FMT, Nfc);

373:   PetscCall(DMSwarmCellDMGetDM(celldm, &dm));
374:   PetscCall(DMGetLocalBoundingBox(dm, lmin, lmax));
375:   PetscCall(DMGetCoordinateDim(dm, &bs));

377:   for (b = 0; b < bs; b++) {
378:     if (npoints[b] > 1) {
379:       dx[b] = (max[b] - min[b]) / ((PetscReal)(npoints[b] - 1));
380:     } else {
381:       dx[b] = 0.0;
382:     }
383:     _npoints[b] = npoints[b];
384:   }

386:   /* determine number of points living in the bounding box */
387:   n_estimate = 0;
388:   for (k = 0; k < _npoints[2]; k++) {
389:     for (j = 0; j < _npoints[1]; j++) {
390:       for (i = 0; i < _npoints[0]; i++) {
391:         PetscReal xp[] = {0.0, 0.0, 0.0};
392:         PetscInt  ijk[3];
393:         PetscBool point_inside = PETSC_TRUE;

395:         ijk[0] = i;
396:         ijk[1] = j;
397:         ijk[2] = k;
398:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
399:         for (b = 0; b < bs; b++) {
400:           if (xp[b] < lmin[b]) point_inside = PETSC_FALSE;
401:           if (xp[b] > lmax[b]) point_inside = PETSC_FALSE;
402:         }
403:         if (point_inside) n_estimate++;
404:       }
405:     }
406:   }

408:   /* create candidate list */
409:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
410:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
411:   PetscCall(VecSetBlockSize(pos, bs));
412:   PetscCall(VecSetFromOptions(pos));
413:   PetscCall(VecGetArray(pos, &_pos));

415:   n_estimate = 0;
416:   for (k = 0; k < _npoints[2]; k++) {
417:     for (j = 0; j < _npoints[1]; j++) {
418:       for (i = 0; i < _npoints[0]; i++) {
419:         PetscReal xp[] = {0.0, 0.0, 0.0};
420:         PetscInt  ijk[3];
421:         PetscBool point_inside = PETSC_TRUE;

423:         ijk[0] = i;
424:         ijk[1] = j;
425:         ijk[2] = k;
426:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
427:         for (b = 0; b < bs; b++) {
428:           if (xp[b] < lmin[b]) point_inside = PETSC_FALSE;
429:           if (xp[b] > lmax[b]) point_inside = PETSC_FALSE;
430:         }
431:         if (point_inside) {
432:           for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = xp[b];
433:           n_estimate++;
434:         }
435:       }
436:     }
437:   }
438:   PetscCall(VecRestoreArray(pos, &_pos));

440:   /* locate points */
441:   PetscCall(DMLocatePoints(dm, pos, DM_POINTLOCATION_NONE, &sfcell));
442:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
443:   n_found = 0;
444:   for (p = 0; p < n_estimate; p++) {
445:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
446:   }

448:   /* adjust size */
449:   if (mode == ADD_VALUES) {
450:     PetscCall(DMSwarmGetLocalSize(sw, &n_curr));
451:     n_new_est = n_curr + n_found;
452:     PetscCall(DMSwarmSetLocalSizes(sw, n_new_est, -1));
453:   }
454:   if (mode == INSERT_VALUES) {
455:     n_curr    = 0;
456:     n_new_est = n_found;
457:     PetscCall(DMSwarmSetLocalSizes(sw, n_new_est, -1));
458:   }

460:   /* initialize new coords, cell owners, pid */
461:   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellid));
462:   PetscCall(VecGetArrayRead(pos, &_coor));
463:   PetscCall(DMSwarmGetField(sw, coordFields[0], NULL, NULL, (void **)&swarm_coor));
464:   PetscCall(DMSwarmGetField(sw, cellid, NULL, NULL, (void **)&swarm_cellid));
465:   n_found = 0;
466:   for (p = 0; p < n_estimate; p++) {
467:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
468:       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
469:       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
470:       n_found++;
471:     }
472:   }
473:   PetscCall(DMSwarmRestoreField(sw, cellid, NULL, NULL, (void **)&swarm_cellid));
474:   PetscCall(DMSwarmRestoreField(sw, coordFields[0], NULL, NULL, (void **)&swarm_coor));
475:   PetscCall(VecRestoreArrayRead(pos, &_coor));

477:   PetscCall(PetscSFDestroy(&sfcell));
478:   PetscCall(VecDestroy(&pos));
479:   PetscFunctionReturn(PETSC_SUCCESS);
480: }

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

485:   Collective

487:   Input Parameters:
488: + sw        - the `DMSWARM`
489: . npoints   - the number of points to insert
490: . coor      - the coordinate values
491: . 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
492: - mode      - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)

494:   Level: beginner

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

501: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`, `DMSwarmSetPointsUniformCoordinates()`
502: @*/
503: PetscErrorCode DMSwarmSetPointCoordinates(DM sw, PetscInt npoints, PetscReal coor[], PetscBool redundant, InsertMode mode)
504: {
505:   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
506:   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
507:   PetscInt           i, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
508:   Vec                coorlocal;
509:   const PetscScalar *_coor;
510:   DMSwarmCellDM      celldm;
511:   DM                 dm;
512:   Vec                pos;
513:   PetscScalar       *_pos;
514:   PetscReal         *swarm_coor;
515:   PetscInt          *swarm_cellid;
516:   PetscSF            sfcell = NULL;
517:   const PetscSFNode *LA_sfcell;
518:   PetscReal         *my_coor;
519:   PetscInt           my_npoints, Nfc;
520:   PetscMPIInt        rank;
521:   MPI_Comm           comm;
522:   const char       **coordFields, *cellid;

524:   PetscFunctionBegin;
525:   DMSWARMPICVALID(sw);
526:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
527:   PetscCallMPI(MPI_Comm_rank(comm, &rank));

529:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
530:   PetscCall(DMSwarmCellDMGetCoordinateFields(celldm, &Nfc, &coordFields));
531:   PetscCheck(Nfc == 1, PetscObjectComm((PetscObject)sw), PETSC_ERR_SUP, "We only support a single coordinate field right now, not %" PetscInt_FMT, Nfc);

533:   PetscCall(DMSwarmCellDMGetDM(celldm, &dm));
534:   PetscCall(DMGetCoordinatesLocal(dm, &coorlocal));
535:   PetscCall(VecGetSize(coorlocal, &N));
536:   PetscCall(VecGetBlockSize(coorlocal, &bs));
537:   N = N / bs;
538:   PetscCall(VecGetArrayRead(coorlocal, &_coor));
539:   for (i = 0; i < N; i++) {
540:     for (b = 0; b < bs; b++) {
541:       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
542:       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
543:     }
544:   }
545:   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));

547:   /* broadcast points from rank 0 if requested */
548:   if (redundant) {
549:     PetscMPIInt imy;

551:     my_npoints = npoints;
552:     PetscCallMPI(MPI_Bcast(&my_npoints, 1, MPIU_INT, 0, comm));

554:     if (rank > 0) { /* allocate space */
555:       PetscCall(PetscMalloc1(bs * my_npoints, &my_coor));
556:     } else {
557:       my_coor = coor;
558:     }
559:     PetscCall(PetscMPIIntCast(bs * my_npoints, &imy));
560:     PetscCallMPI(MPI_Bcast(my_coor, imy, MPIU_REAL, 0, comm));
561:   } else {
562:     my_npoints = npoints;
563:     my_coor    = coor;
564:   }

566:   /* determine the number of points living in the bounding box */
567:   n_estimate = 0;
568:   for (i = 0; i < my_npoints; i++) {
569:     PetscBool point_inside = PETSC_TRUE;

571:     for (b = 0; b < bs; b++) {
572:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
573:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
574:     }
575:     if (point_inside) n_estimate++;
576:   }

578:   /* create candidate list */
579:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
580:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
581:   PetscCall(VecSetBlockSize(pos, bs));
582:   PetscCall(VecSetFromOptions(pos));
583:   PetscCall(VecGetArray(pos, &_pos));

585:   n_estimate = 0;
586:   for (i = 0; i < my_npoints; i++) {
587:     PetscBool point_inside = PETSC_TRUE;

589:     for (b = 0; b < bs; b++) {
590:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
591:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
592:     }
593:     if (point_inside) {
594:       for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = my_coor[bs * i + b];
595:       n_estimate++;
596:     }
597:   }
598:   PetscCall(VecRestoreArray(pos, &_pos));

600:   /* locate points */
601:   PetscCall(DMLocatePoints(dm, pos, DM_POINTLOCATION_NONE, &sfcell));

603:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
604:   n_found = 0;
605:   for (p = 0; p < n_estimate; p++) {
606:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
607:   }

609:   /* adjust size */
610:   if (mode == ADD_VALUES) {
611:     PetscCall(DMSwarmGetLocalSize(sw, &n_curr));
612:     n_new_est = n_curr + n_found;
613:     PetscCall(DMSwarmSetLocalSizes(sw, n_new_est, -1));
614:   }
615:   if (mode == INSERT_VALUES) {
616:     n_curr    = 0;
617:     n_new_est = n_found;
618:     PetscCall(DMSwarmSetLocalSizes(sw, n_new_est, -1));
619:   }

621:   /* initialize new coords, cell owners, pid */
622:   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellid));
623:   PetscCall(VecGetArrayRead(pos, &_coor));
624:   PetscCall(DMSwarmGetField(sw, coordFields[0], NULL, NULL, (void **)&swarm_coor));
625:   PetscCall(DMSwarmGetField(sw, cellid, NULL, NULL, (void **)&swarm_cellid));
626:   n_found = 0;
627:   for (p = 0; p < n_estimate; p++) {
628:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
629:       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
630:       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
631:       n_found++;
632:     }
633:   }
634:   PetscCall(DMSwarmRestoreField(sw, cellid, NULL, NULL, (void **)&swarm_cellid));
635:   PetscCall(DMSwarmRestoreField(sw, coordFields[0], NULL, NULL, (void **)&swarm_coor));
636:   PetscCall(VecRestoreArrayRead(pos, &_coor));

638:   if (redundant) {
639:     if (rank > 0) PetscCall(PetscFree(my_coor));
640:   }
641:   PetscCall(PetscSFDestroy(&sfcell));
642:   PetscCall(VecDestroy(&pos));
643:   PetscFunctionReturn(PETSC_SUCCESS);
644: }

646: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
647: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);

649: /*@
650:   DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell

652:   Not Collective

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

659:   Level: beginner

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

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

666:   When using a `DMPLEX` the following case are supported\:
667: .vb
668:    (i) DMSWARMPIC_LAYOUT_REGULAR: 2D (triangle),
669:    (ii) DMSWARMPIC_LAYOUT_GAUSS: 2D and 3D provided the cell is a tri/tet or a quad/hex,
670:    (iii) DMSWARMPIC_LAYOUT_SUBDIVISION: 2D and 3D for quad/hex and 2D tri.
671: .ve

673: .seealso: `DMSWARM`, `DMSwarmPICLayoutType`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
674: @*/
675: PetscErrorCode DMSwarmInsertPointsUsingCellDM(DM dm, DMSwarmPICLayoutType layout_type, PetscInt fill_param)
676: {
677:   DM        celldm;
678:   PetscBool isDA, isPLEX;

680:   PetscFunctionBegin;
681:   DMSWARMPICVALID(dm);
682:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
683:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
684:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
685:   if (isDA) {
686:     PetscCall(private_DMSwarmInsertPointsUsingCellDM_DA(dm, celldm, layout_type, fill_param));
687:   } else if (isPLEX) {
688:     PetscCall(private_DMSwarmInsertPointsUsingCellDM_PLEX(dm, celldm, layout_type, fill_param));
689:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
690:   PetscFunctionReturn(PETSC_SUCCESS);
691: }

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

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

698:   Not Collective

700:   Input Parameters:
701: + dm      - the `DMSWARM`
702: . npoints - the number of points to insert in each cell
703: - xi      - the coordinates (defined in the local coordinate system for each cell) to insert

705:   Level: beginner

707:   Notes:
708:   The method will reset any previous defined points within the `DMSWARM`.
709:   Only supported for `DMPLEX`. If you are using a `DMDA` it is recommended to either use
710:   `DMSwarmInsertPointsUsingCellDM()`, or extract and set the coordinates yourself the following code
711: .vb
712:     PetscReal  *coor;
713:     const char *coordname;
714:     DMSwarmGetCoordinateField(dm, &coordname);
715:     DMSwarmGetField(dm,coordname,NULL,NULL,(void**)&coor);
716:     // user code to define the coordinates here
717:     DMSwarmRestoreField(dm,coordname,NULL,NULL,(void**)&coor);
718: .ve

720: .seealso: `DMSWARM`, `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
721: @*/
722: PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
723: {
724:   DM        celldm;
725:   PetscBool isDA, isPLEX;

727:   PetscFunctionBegin;
728:   DMSWARMPICVALID(dm);
729:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
730:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
731:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
732:   PetscCheck(!isDA, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMPLEX. Recommended you use DMSwarmInsertPointsUsingCellDM()");
733:   PetscCheck(isPLEX, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
734:   PetscCall(private_DMSwarmSetPointCoordinatesCellwise_PLEX(dm, celldm, npoints, xi));
735:   PetscFunctionReturn(PETSC_SUCCESS);
736: }

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

741:   Not Collective

743:   Input Parameter:
744: . sw - the `DMSWARM`

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

750:   Level: beginner

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

755: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
756: @*/
757: PetscErrorCode DMSwarmCreatePointPerCellCount(DM sw, PetscInt *ncells, PetscInt **count)
758: {
759:   DMSwarmCellDM celldm;
760:   PetscBool     isvalid;
761:   PetscInt      nel;
762:   PetscInt     *sum;
763:   const char   *cellid;

765:   PetscFunctionBegin;
766:   PetscCall(DMSwarmSortGetIsValid(sw, &isvalid));
767:   nel = 0;
768:   if (isvalid) {
769:     PetscInt e;

771:     PetscCall(DMSwarmSortGetSizes(sw, &nel, NULL));

773:     PetscCall(PetscMalloc1(nel, &sum));
774:     for (e = 0; e < nel; e++) PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, e, &sum[e]));
775:   } else {
776:     DM        dm;
777:     PetscBool isda, isplex, isshell;
778:     PetscInt  p, npoints;
779:     PetscInt *swarm_cellid;

781:     /* get the number of cells */
782:     PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
783:     PetscCall(DMSwarmCellDMGetDM(celldm, &dm));
784:     PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMDA, &isda));
785:     PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMPLEX, &isplex));
786:     PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMSHELL, &isshell));
787:     if (isda) {
788:       PetscInt        _nel, _npe;
789:       const PetscInt *_element;

791:       PetscCall(DMDAGetElements(dm, &_nel, &_npe, &_element));
792:       nel = _nel;
793:       PetscCall(DMDARestoreElements(dm, &_nel, &_npe, &_element));
794:     } else if (isplex) {
795:       PetscInt ps, pe;

797:       PetscCall(DMPlexGetHeightStratum(dm, 0, &ps, &pe));
798:       nel = pe - ps;
799:     } else if (isshell) {
800:       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);

802:       PetscCall(PetscObjectQueryFunction((PetscObject)dm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
803:       if (method_DMShellGetNumberOfCells) {
804:         PetscCall(method_DMShellGetNumberOfCells(dm, &nel));
805:       } else
806:         SETERRQ(PetscObjectComm((PetscObject)sw), 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);");
807:     } else SETERRQ(PetscObjectComm((PetscObject)sw), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");

809:     PetscCall(PetscMalloc1(nel, &sum));
810:     PetscCall(PetscArrayzero(sum, nel));
811:     PetscCall(DMSwarmGetLocalSize(sw, &npoints));
812:     PetscCall(DMSwarmCellDMGetCellID(celldm, &cellid));
813:     PetscCall(DMSwarmGetField(sw, cellid, NULL, NULL, (void **)&swarm_cellid));
814:     for (p = 0; p < npoints; p++) {
815:       if (swarm_cellid[p] != DMLOCATEPOINT_POINT_NOT_FOUND) sum[swarm_cellid[p]]++;
816:     }
817:     PetscCall(DMSwarmRestoreField(sw, cellid, NULL, NULL, (void **)&swarm_cellid));
818:   }
819:   if (ncells) *ncells = nel;
820:   *count = sum;
821:   PetscFunctionReturn(PETSC_SUCCESS);
822: }

824: /*@
825:   DMSwarmGetNumSpecies - Get the number of particle species

827:   Not Collective

829:   Input Parameter:
830: . sw - the `DMSWARM`

832:   Output Parameters:
833: . Ns - the number of species

835:   Level: intermediate

837: .seealso: `DMSWARM`, `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
838: @*/
839: PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
840: {
841:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

843:   PetscFunctionBegin;
844:   *Ns = swarm->Ns;
845:   PetscFunctionReturn(PETSC_SUCCESS);
846: }

848: /*@
849:   DMSwarmSetNumSpecies - Set the number of particle species

851:   Not Collective

853:   Input Parameters:
854: + sw - the `DMSWARM`
855: - Ns - the number of species

857:   Level: intermediate

859: .seealso: `DMSWARM`, `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
860: @*/
861: PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
862: {
863:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

865:   PetscFunctionBegin;
866:   swarm->Ns = Ns;
867:   PetscFunctionReturn(PETSC_SUCCESS);
868: }

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

873:   Not Collective

875:   Input Parameter:
876: . sw - the `DMSWARM`

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

881:   Level: intermediate

883: .seealso: `DMSWARM`, `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`, `PetscSimplePointFn`
884: @*/
885: PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFn **coordFunc)
886: {
887:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

889:   PetscFunctionBegin;
891:   PetscAssertPointer(coordFunc, 2);
892:   *coordFunc = swarm->coordFunc;
893:   PetscFunctionReturn(PETSC_SUCCESS);
894: }

896: /*@C
897:   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions

899:   Not Collective

901:   Input Parameters:
902: + sw        - the `DMSWARM`
903: - coordFunc - the function setting initial particle positions, see `PetscSimplePointFn` for the calling sequence

905:   Level: intermediate

907: .seealso: `DMSWARM`, `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`, `PetscSimplePointFn`
908: @*/
909: PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFn *coordFunc)
910: {
911:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

913:   PetscFunctionBegin;
916:   swarm->coordFunc = coordFunc;
917:   PetscFunctionReturn(PETSC_SUCCESS);
918: }

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

923:   Not Collective

925:   Input Parameter:
926: . sw - the `DMSWARM`

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

931:   Level: intermediate

933: .seealso: `DMSWARM`, `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`, `PetscSimplePointFn`
934: @*/
935: PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFn **velFunc)
936: {
937:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

939:   PetscFunctionBegin;
941:   PetscAssertPointer(velFunc, 2);
942:   *velFunc = swarm->velFunc;
943:   PetscFunctionReturn(PETSC_SUCCESS);
944: }

946: /*@C
947:   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities

949:   Not Collective

951:   Input Parameters:
952: + sw      - the `DMSWARM`
953: - velFunc - the function setting initial particle velocities, see `PetscSimplePointFn` for the calling sequence

955:   Level: intermediate

957: .seealso: `DMSWARM`, `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`, `PetscSimplePointFn`
958: @*/
959: PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFn *velFunc)
960: {
961:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

963:   PetscFunctionBegin;
966:   swarm->velFunc = velFunc;
967:   PetscFunctionReturn(PETSC_SUCCESS);
968: }

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

973:   Not Collective

975:   Input Parameters:
976: + sw      - The `DMSWARM`
977: . N       - The target number of particles
978: - density - The density field for the particle layout, normalized to unity

980:   Level: advanced

982:   Note:
983:   One particle will be created for each species.

985: .seealso: `DMSWARM`, `DMSwarmComputeLocalSizeFromOptions()`
986: @*/
987: PetscErrorCode DMSwarmComputeLocalSize(DM sw, PetscInt N, PetscProbFn *density)
988: {
989:   DM               dm;
990:   DMSwarmCellDM    celldm;
991:   PetscQuadrature  quad;
992:   const PetscReal *xq, *wq;
993:   PetscReal       *n_int;
994:   PetscInt        *npc_s, *swarm_cellid, Ni;
995:   PetscReal        gmin[3], gmax[3], xi0[3];
996:   PetscInt         Ns, cStart, cEnd, c, dim, d, Nq, q, Np = 0, p, s;
997:   PetscBool        simplex;
998:   const char      *cellid;

1000:   PetscFunctionBegin;
1001:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1002:   PetscCall(DMSwarmGetCellDM(sw, &dm));
1003:   PetscCall(DMGetDimension(dm, &dim));
1004:   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
1005:   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1006:   PetscCall(DMPlexIsSimplex(dm, &simplex));
1007:   PetscCall(DMGetCoordinatesLocalSetUp(dm));
1008:   if (simplex) PetscCall(PetscDTStroudConicalQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
1009:   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
1010:   PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, &xq, &wq));
1011:   PetscCall(PetscCalloc2(Ns, &n_int, (cEnd - cStart) * Ns, &npc_s));
1012:   /* Integrate the density function to get the number of particles in each cell */
1013:   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
1014:   for (c = 0; c < cEnd - cStart; ++c) {
1015:     const PetscInt cell = c + cStart;
1016:     PetscReal      v0[3], J[9], invJ[9], detJ, detJp = 2. / (gmax[0] - gmin[0]), xr[3], den;

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

1026:       for (s = 0; s < Ns; ++s) {
1027:         PetscCall(density(xr, NULL, &den));
1028:         n_int[s] += (detJp * den) * (detJ * wq[q]) / (PetscReal)Ns;
1029:       }
1030:     }
1031:     for (s = 0; s < Ns; ++s) {
1032:       Ni = N;
1033:       npc_s[c * Ns + s] += (PetscInt)(Ni * n_int[s] + 0.5); // TODO Wish we wrapped round()
1034:       Np += npc_s[c * Ns + s];
1035:     }
1036:   }
1037:   PetscCall(PetscQuadratureDestroy(&quad));
1038:   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
1039:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
1040:   PetscCall(DMSwarmCellDMGetCellID(celldm, &cellid));
1041:   PetscCall(DMSwarmGetField(sw, cellid, NULL, NULL, (void **)&swarm_cellid));
1042:   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
1043:     for (s = 0; s < Ns; ++s) {
1044:       for (q = 0; q < npc_s[c * Ns + s]; ++q, ++p) swarm_cellid[p] = c;
1045:     }
1046:   }
1047:   PetscCall(DMSwarmRestoreField(sw, cellid, NULL, NULL, (void **)&swarm_cellid));
1048:   PetscCall(PetscFree2(n_int, npc_s));
1049:   PetscFunctionReturn(PETSC_SUCCESS);
1050: }

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

1055:   Not Collective

1057:   Input Parameter:
1058: . sw - The `DMSWARM`

1060:   Level: advanced

1062: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`
1063: @*/
1064: PetscErrorCode DMSwarmComputeLocalSizeFromOptions(DM sw)
1065: {
1066:   PetscProbFn *pdf;
1067:   const char  *prefix;
1068:   char         funcname[PETSC_MAX_PATH_LEN];
1069:   PetscInt    *N, Ns, dim, n;
1070:   PetscBool    flg;
1071:   PetscMPIInt  size, rank;

1073:   PetscFunctionBegin;
1074:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
1075:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
1076:   PetscCall(PetscCalloc1(size, &N));
1077:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
1078:   n = size;
1079:   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
1080:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1081:   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
1082:   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
1083:   PetscCall(PetscOptionsString("-dm_swarm_coordinate_function", "Function to determine particle coordinates", "DMSwarmSetCoordinateFunction", funcname, funcname, sizeof(funcname), &flg));
1084:   PetscOptionsEnd();
1085:   if (flg) {
1086:     PetscSimplePointFn *coordFunc;

1088:     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1089:     PetscCall(PetscDLSym(NULL, funcname, (void **)&coordFunc));
1090:     PetscCheck(coordFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
1091:     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1092:     PetscCall(DMSwarmSetLocalSizes(sw, N[rank] * Ns, 0));
1093:     PetscCall(DMSwarmSetCoordinateFunction(sw, coordFunc));
1094:   } else {
1095:     PetscCall(DMGetDimension(sw, &dim));
1096:     PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
1097:     PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_coordinate_density", &pdf, NULL, NULL));
1098:     PetscCall(DMSwarmComputeLocalSize(sw, N[rank], pdf));
1099:   }
1100:   PetscCall(PetscFree(N));
1101:   PetscFunctionReturn(PETSC_SUCCESS);
1102: }

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

1107:   Not Collective

1109:   Input Parameter:
1110: . sw - The `DMSWARM`

1112:   Level: advanced

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

1117: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeVelocities()`
1118: @*/
1119: PetscErrorCode DMSwarmInitializeCoordinates(DM sw)
1120: {
1121:   DMSwarmCellDM       celldm;
1122:   PetscSimplePointFn *coordFunc;
1123:   PetscScalar        *weight;
1124:   PetscReal          *x;
1125:   PetscInt           *species;
1126:   void               *ctx;
1127:   PetscBool           removePoints = PETSC_TRUE;
1128:   PetscDataType       dtype;
1129:   PetscInt            Nfc, Np, p, Ns, dim, d, bs;
1130:   const char        **coordFields;

1132:   PetscFunctionBeginUser;
1133:   PetscCall(DMGetDimension(sw, &dim));
1134:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1135:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1136:   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));

1138:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
1139:   PetscCall(DMSwarmCellDMGetCoordinateFields(celldm, &Nfc, &coordFields));
1140:   PetscCheck(Nfc == 1, PetscObjectComm((PetscObject)sw), PETSC_ERR_SUP, "We only support a single coordinate field right now, not %" PetscInt_FMT, Nfc);

1142:   PetscCall(DMSwarmGetField(sw, coordFields[0], &bs, &dtype, (void **)&x));
1143:   PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight));
1144:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1145:   if (coordFunc) {
1146:     PetscCall(DMGetApplicationContext(sw, &ctx));
1147:     for (p = 0; p < Np; ++p) {
1148:       PetscScalar X[3];

1150:       PetscCall((*coordFunc)(dim, 0., NULL, p, X, ctx));
1151:       for (d = 0; d < dim; ++d) x[p * dim + d] = PetscRealPart(X[d]);
1152:       weight[p]  = 1.0;
1153:       species[p] = p % Ns;
1154:     }
1155:   } else {
1156:     DM          dm;
1157:     PetscRandom rnd;
1158:     PetscReal   xi0[3];
1159:     PetscInt    cStart, cEnd, c;

1161:     PetscCall(DMSwarmGetCellDM(sw, &dm));
1162:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1163:     PetscCall(DMGetApplicationContext(sw, &ctx));

1165:     /* Set particle position randomly in cell, set weights to 1 */
1166:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd));
1167:     PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
1168:     PetscCall(PetscRandomSetFromOptions(rnd));
1169:     PetscCall(DMSwarmSortGetAccess(sw));
1170:     for (d = 0; d < dim; ++d) xi0[d] = -1.0;
1171:     for (c = cStart; c < cEnd; ++c) {
1172:       PetscReal v0[3], J[9], invJ[9], detJ;
1173:       PetscInt *pidx, Npc, q;

1175:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1176:       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ));
1177:       for (q = 0; q < Npc; ++q) {
1178:         const PetscInt p = pidx[q];
1179:         PetscReal      xref[3];

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

1184:         weight[p]  = 1.0 / Np;
1185:         species[p] = p % Ns;
1186:       }
1187:       PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx));
1188:     }
1189:     PetscCall(PetscRandomDestroy(&rnd));
1190:     PetscCall(DMSwarmSortRestoreAccess(sw));
1191:   }
1192:   PetscCall(DMSwarmRestoreField(sw, coordFields[0], NULL, NULL, (void **)&x));
1193:   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
1194:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));

1196:   PetscCall(DMSwarmMigrate(sw, removePoints));
1197:   PetscCall(DMLocalizeCoordinates(sw));
1198:   PetscFunctionReturn(PETSC_SUCCESS);
1199: }

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

1204:   Collective

1206:   Input Parameters:
1207: + sw      - The `DMSWARM` object
1208: . sampler - A function which uniformly samples the velocity PDF
1209: - v0      - The velocity scale for nondimensionalization for each species

1211:   Level: advanced

1213:   Note:
1214:   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.

1216: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocitiesFromOptions()`
1217: @*/
1218: PetscErrorCode DMSwarmInitializeVelocities(DM sw, PetscProbFn *sampler, const PetscReal v0[])
1219: {
1220:   PetscSimplePointFn *velFunc;
1221:   PetscReal          *v;
1222:   PetscInt           *species;
1223:   void               *ctx;
1224:   PetscInt            dim, Np, p;

1226:   PetscFunctionBegin;
1227:   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));

1229:   PetscCall(DMGetDimension(sw, &dim));
1230:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1231:   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
1232:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1233:   if (v0[0] == 0.) {
1234:     PetscCall(PetscArrayzero(v, Np * dim));
1235:   } else if (velFunc) {
1236:     PetscCall(DMGetApplicationContext(sw, &ctx));
1237:     for (p = 0; p < Np; ++p) {
1238:       PetscInt    s = species[p], d;
1239:       PetscScalar vel[3];

1241:       PetscCall((*velFunc)(dim, 0., NULL, p, vel, ctx));
1242:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
1243:     }
1244:   } else {
1245:     PetscRandom rnd;

1247:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
1248:     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
1249:     PetscCall(PetscRandomSetFromOptions(rnd));

1251:     for (p = 0; p < Np; ++p) {
1252:       PetscInt  s = species[p], d;
1253:       PetscReal a[3], vel[3];

1255:       for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &a[d]));
1256:       PetscCall(sampler(a, NULL, vel));
1257:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * vel[d];
1258:     }
1259:     PetscCall(PetscRandomDestroy(&rnd));
1260:   }
1261:   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1262:   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1263:   PetscFunctionReturn(PETSC_SUCCESS);
1264: }

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

1269:   Collective

1271:   Input Parameters:
1272: + sw - The `DMSWARM` object
1273: - v0 - The velocity scale for nondimensionalization for each species

1275:   Level: advanced

1277: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
1278: @*/
1279: PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
1280: {
1281:   PetscProbFn *sampler;
1282:   PetscInt     dim;
1283:   const char  *prefix;
1284:   char         funcname[PETSC_MAX_PATH_LEN];
1285:   PetscBool    flg;

1287:   PetscFunctionBegin;
1288:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
1289:   PetscCall(PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg));
1290:   PetscOptionsEnd();
1291:   if (flg) {
1292:     PetscSimplePointFn *velFunc;

1294:     PetscCall(PetscDLSym(NULL, funcname, (void **)&velFunc));
1295:     PetscCheck(velFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
1296:     PetscCall(DMSwarmSetVelocityFunction(sw, velFunc));
1297:   }
1298:   PetscCall(DMGetDimension(sw, &dim));
1299:   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
1300:   PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_velocity_density", NULL, NULL, &sampler));
1301:   PetscCall(DMSwarmInitializeVelocities(sw, sampler, v0));
1302:   PetscFunctionReturn(PETSC_SUCCESS);
1303: }

1305: // The input vector U is assumed to be from a PetscFE. The Swarm fields are input as auxiliary values.
1306: PetscErrorCode DMProjectFieldLocal_Swarm(DM dm, PetscReal time, Vec U, PetscPointFn **funcs, InsertMode mode, Vec X)
1307: {
1308:   MPI_Comm         comm;
1309:   DM               dmIn;
1310:   PetscDS          ds;
1311:   PetscTabulation *T;
1312:   DMSwarmCellDM    celldm;
1313:   PetscScalar     *a, *val, *u, *u_x;
1314:   PetscFEGeom      fegeom;
1315:   PetscReal       *xi, *v0, *J, *invJ, detJ = 1.0, v0ref[3] = {-1.0, -1.0, -1.0};
1316:   PetscInt         dim, dE, Np, n, Nf, Nfc, Nfu, cStart, cEnd, maxC = 0, totbs = 0;
1317:   const char     **coordFields, **fields;
1318:   PetscReal      **coordVals, **vals;
1319:   PetscInt        *cbs, *bs, *uOff, *uOff_x;

1321:   PetscFunctionBegin;
1322:   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
1323:   PetscCall(VecGetDM(U, &dmIn));
1324:   PetscCall(DMGetDimension(dmIn, &dim));
1325:   PetscCall(DMGetCoordinateDim(dmIn, &dE));
1326:   PetscCall(DMGetDS(dmIn, &ds));
1327:   PetscCall(PetscDSGetNumFields(ds, &Nfu));
1328:   PetscCall(PetscDSGetComponentOffsets(ds, &uOff));
1329:   PetscCall(PetscDSGetComponentDerivativeOffsets(ds, &uOff_x));
1330:   PetscCall(PetscDSGetTabulation(ds, &T));
1331:   PetscCall(PetscDSGetEvaluationArrays(ds, &u, NULL, &u_x));
1332:   PetscCall(PetscMalloc3(dim, &v0, dim * dim, &J, dim * dim, &invJ));
1333:   PetscCall(DMPlexGetHeightStratum(dmIn, 0, &cStart, &cEnd));

1335:   fegeom.dim      = dim;
1336:   fegeom.dimEmbed = dE;
1337:   fegeom.v        = v0;
1338:   fegeom.xi       = v0ref;
1339:   fegeom.J        = J;
1340:   fegeom.invJ     = invJ;
1341:   fegeom.detJ     = &detJ;

1343:   PetscCall(DMSwarmGetLocalSize(dm, &Np));
1344:   PetscCall(VecGetLocalSize(X, &n));
1345:   PetscCheck(n == Np, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Output vector local size %" PetscInt_FMT " != %" PetscInt_FMT " number of local particles", n, Np);
1346:   PetscCall(DMSwarmGetCellDMActive(dm, &celldm));
1347:   PetscCall(DMSwarmCellDMGetCoordinateFields(celldm, &Nfc, &coordFields));
1348:   PetscCall(DMSwarmCellDMGetFields(celldm, &Nf, &fields));

1350:   PetscCall(PetscMalloc2(Nfc, &coordVals, Nfc, &cbs));
1351:   for (PetscInt i = 0; i < Nfc; ++i) PetscCall(DMSwarmGetField(dm, coordFields[i], &cbs[i], NULL, (void **)&coordVals[i]));
1352:   PetscCall(PetscMalloc2(Nf, &vals, Nfc, &bs));
1353:   for (PetscInt i = 0; i < Nf; ++i) {
1354:     PetscCall(DMSwarmGetField(dm, fields[i], &bs[i], NULL, (void **)&vals[i]));
1355:     totbs += bs[i];
1356:   }

1358:   PetscCall(DMSwarmSortGetAccess(dm));
1359:   for (PetscInt cell = cStart; cell < cEnd; ++cell) {
1360:     PetscInt *pindices, Npc;

1362:     PetscCall(DMSwarmSortGetPointsPerCell(dm, cell, &Npc, &pindices));
1363:     maxC = PetscMax(maxC, Npc);
1364:     PetscCall(DMSwarmSortRestorePointsPerCell(dm, cell, &Npc, &pindices));
1365:   }
1366:   PetscCall(PetscMalloc3(maxC * dim, &xi, maxC * totbs, &val, Nfu, &T));
1367:   PetscCall(VecGetArray(X, &a));
1368:   {
1369:     for (PetscInt cell = cStart; cell < cEnd; ++cell) {
1370:       PetscInt *pindices, Npc;

1372:       // TODO: Use DMField instead of assuming affine
1373:       PetscCall(DMPlexComputeCellGeometryFEM(dmIn, cell, NULL, v0, J, invJ, &detJ));
1374:       PetscCall(DMSwarmSortGetPointsPerCell(dm, cell, &Npc, &pindices));

1376:       PetscScalar *closure = NULL;
1377:       PetscInt     Ncl;

1379:       // Get fields from input vector and auxiliary fields from swarm
1380:       for (PetscInt p = 0; p < Npc; ++p) {
1381:         PetscReal xr[8];
1382:         PetscInt  off;

1384:         off = 0;
1385:         for (PetscInt i = 0; i < Nfc; ++i) {
1386:           for (PetscInt b = 0; b < cbs[i]; ++b, ++off) xr[off] = coordVals[i][pindices[p] * cbs[i] + b];
1387:         }
1388:         PetscCheck(off == dim, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "The total block size of coordinates is %" PetscInt_FMT " != %" PetscInt_FMT " the DM coordinate dimension", off, dim);
1389:         CoordinatesRealToRef(dE, dim, fegeom.xi, fegeom.v, fegeom.invJ, xr, &xi[p * dim]);
1390:         off = 0;
1391:         for (PetscInt i = 0; i < Nf; ++i) {
1392:           for (PetscInt b = 0; b < bs[i]; ++b, ++off) val[p * totbs + off] = vals[i][pindices[p] * bs[i] + b];
1393:         }
1394:         PetscCheck(off == totbs, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "The total block size of swarm fields is %" PetscInt_FMT " != %" PetscInt_FMT " the computed total block size", off, totbs);
1395:       }
1396:       PetscCall(DMPlexVecGetClosure(dmIn, NULL, U, cell, &Ncl, &closure));
1397:       for (PetscInt field = 0; field < Nfu; ++field) {
1398:         PetscFE fe;

1400:         PetscCall(PetscDSGetDiscretization(ds, field, (PetscObject *)&fe));
1401:         PetscCall(PetscFECreateTabulation(fe, 1, Npc, xi, 1, &T[field]));
1402:       }
1403:       for (PetscInt p = 0; p < Npc; ++p) {
1404:         // Get fields from input vector
1405:         PetscCall(PetscFEEvaluateFieldJets_Internal(ds, Nfu, 0, p, T, &fegeom, closure, NULL, u, u_x, NULL));
1406:         (*funcs[0])(dim, 1, 1, uOff, uOff_x, u, NULL, u_x, bs, NULL, &val[p * totbs], NULL, NULL, time, &xi[p * dim], 0, NULL, &a[pindices[p]]);
1407:       }
1408:       PetscCall(DMSwarmSortRestorePointsPerCell(dm, cell, &Npc, &pindices));
1409:       PetscCall(DMPlexVecRestoreClosure(dmIn, NULL, U, cell, &Ncl, &closure));
1410:       for (PetscInt field = 0; field < Nfu; ++field) PetscCall(PetscTabulationDestroy(&T[field]));
1411:     }
1412:   }
1413:   for (PetscInt i = 0; i < Nfc; ++i) PetscCall(DMSwarmRestoreField(dm, coordFields[i], &cbs[i], NULL, (void **)&coordVals[i]));
1414:   for (PetscInt i = 0; i < Nf; ++i) PetscCall(DMSwarmRestoreField(dm, fields[i], &bs[i], NULL, (void **)&vals[i]));
1415:   PetscCall(VecRestoreArray(X, &a));
1416:   PetscCall(DMSwarmSortRestoreAccess(dm));
1417:   PetscCall(PetscFree3(xi, val, T));
1418:   PetscCall(PetscFree3(v0, J, invJ));
1419:   PetscCall(PetscFree2(coordVals, cbs));
1420:   PetscCall(PetscFree2(vals, bs));
1421:   PetscFunctionReturn(PETSC_SUCCESS);
1422: }