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()`
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

229:   Level: intermediate

231: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`
232: @*/
233: PetscErrorCode DMSwarmCellDMSetSort(DMSwarmCellDM celldm, DMSwarmSort sort)
234: {
235:   PetscFunctionBegin;
237:   PetscAssertPointer(sort, 2);
238:   celldm->sort = sort;
239:   PetscFunctionReturn(PETSC_SUCCESS);
240: }

242: /*@
243:   DMSwarmCellDMGetBlockSize - Returns the total blocksize for the `DM` fields

245:   Not Collective

247:   Input Parameters:
248: + celldm - The `DMSwarmCellDM` object
249: - sw     - The `DMSwarm` object

251:   Output Parameter:
252: . bs - The total block size

254:   Level: intermediate

256: .seealso: `DMSwarmCellDM`, `DM`, `DMSwarmSetCellDM()`
257: @*/
258: PetscErrorCode DMSwarmCellDMGetBlockSize(DMSwarmCellDM celldm, DM sw, PetscInt *bs)
259: {
260:   PetscFunctionBegin;
263:   PetscAssertPointer(bs, 3);
264:   *bs = 0;
265:   for (PetscInt f = 0; f < celldm->Nf; ++f) {
266:     PetscInt fbs;

268:     PetscCall(DMSwarmGetFieldInfo(sw, celldm->dmFields[f], &fbs, NULL));
269:     *bs += fbs;
270:   }
271:   PetscFunctionReturn(PETSC_SUCCESS);
272: }

274: /*@
275:   DMSwarmCellDMCreate - create a `DMSwarmCellDM`

277:   Collective

279:   Input Parameters:
280: + dm          - The background `DM` for the `DMSwarm`
281: . Nf          - The number of swarm fields defined over `dm`
282: . dmFields    - The swarm field names for the `dm` fields
283: . Nfc         - The number of swarm fields to use for coordinates over `dm`
284: - coordFields - The swarm field names for the `dm` coordinate fields

286:   Output Parameter:
287: . celldm - The new `DMSwarmCellDM`

289:   Level: advanced

291: .seealso: `DMSwarmCellDM`, `DMSWARM`, `DMSetType()`
292: @*/
293: PetscErrorCode DMSwarmCellDMCreate(DM dm, PetscInt Nf, const char *dmFields[], PetscInt Nfc, const char *coordFields[], DMSwarmCellDM *celldm)
294: {
295:   DMSwarmCellDM b;
296:   const char   *name;
297:   char          cellid[PETSC_MAX_PATH_LEN];

299:   PetscFunctionBegin;
301:   if (Nf) PetscAssertPointer(dmFields, 3);
302:   if (Nfc) PetscAssertPointer(coordFields, 5);
303:   PetscCall(DMInitializePackage());

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

322: /* Coordinate insertition/addition API */
323: /*@
324:   DMSwarmSetPointsUniformCoordinates - Set point coordinates in a `DMSWARM` on a regular (ijk) grid

326:   Collective

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

335:   Level: beginner

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

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

362:   PetscFunctionBegin;
363:   DMSWARMPICVALID(sw);
364:   PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
365:   PetscCall(DMSwarmCellDMGetCoordinateFields(celldm, &Nfc, &coordFields));
366:   PetscCheck(Nfc == 1, PetscObjectComm((PetscObject)sw), PETSC_ERR_SUP, "We only support a single coordinate field right now, not %" PetscInt_FMT, Nfc);

368:   PetscCall(DMSwarmCellDMGetDM(celldm, &dm));
369:   PetscCall(DMGetLocalBoundingBox(dm, lmin, lmax));
370:   PetscCall(DMGetCoordinateDim(dm, &bs));

372:   for (b = 0; b < bs; b++) {
373:     if (npoints[b] > 1) {
374:       dx[b] = (max[b] - min[b]) / ((PetscReal)(npoints[b] - 1));
375:     } else {
376:       dx[b] = 0.0;
377:     }
378:     _npoints[b] = npoints[b];
379:   }

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

390:         ijk[0] = i;
391:         ijk[1] = j;
392:         ijk[2] = k;
393:         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
394:         for (b = 0; b < bs; b++) {
395:           if (xp[b] < lmin[b]) point_inside = PETSC_FALSE;
396:           if (xp[b] > lmax[b]) point_inside = PETSC_FALSE;
397:         }
398:         if (point_inside) n_estimate++;
399:       }
400:     }
401:   }

403:   /* create candidate list */
404:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
405:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
406:   PetscCall(VecSetBlockSize(pos, bs));
407:   PetscCall(VecSetFromOptions(pos));
408:   PetscCall(VecGetArray(pos, &_pos));

410:   n_estimate = 0;
411:   for (k = 0; k < _npoints[2]; k++) {
412:     for (j = 0; j < _npoints[1]; j++) {
413:       for (i = 0; i < _npoints[0]; i++) {
414:         PetscReal xp[] = {0.0, 0.0, 0.0};
415:         PetscInt  ijk[3];
416:         PetscBool point_inside = PETSC_TRUE;

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

435:   /* locate points */
436:   PetscCall(DMLocatePoints(dm, pos, DM_POINTLOCATION_NONE, &sfcell));
437:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
438:   n_found = 0;
439:   for (p = 0; p < n_estimate; p++) {
440:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
441:   }

443:   /* adjust size */
444:   if (mode == ADD_VALUES) {
445:     PetscCall(DMSwarmGetLocalSize(sw, &n_curr));
446:     n_new_est = n_curr + n_found;
447:     PetscCall(DMSwarmSetLocalSizes(sw, n_new_est, -1));
448:   }
449:   if (mode == INSERT_VALUES) {
450:     n_curr    = 0;
451:     n_new_est = n_found;
452:     PetscCall(DMSwarmSetLocalSizes(sw, n_new_est, -1));
453:   }

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

472:   PetscCall(PetscSFDestroy(&sfcell));
473:   PetscCall(VecDestroy(&pos));
474:   PetscFunctionReturn(PETSC_SUCCESS);
475: }

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

480:   Collective

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

489:   Level: beginner

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

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

519:   PetscFunctionBegin;
520:   DMSWARMPICVALID(sw);
521:   PetscCall(PetscObjectGetComm((PetscObject)sw, &comm));
522:   PetscCallMPI(MPI_Comm_rank(comm, &rank));

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

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

542:   /* broadcast points from rank 0 if requested */
543:   if (redundant) {
544:     PetscMPIInt imy;

546:     my_npoints = npoints;
547:     PetscCallMPI(MPI_Bcast(&my_npoints, 1, MPIU_INT, 0, comm));

549:     if (rank > 0) { /* allocate space */
550:       PetscCall(PetscMalloc1(bs * my_npoints, &my_coor));
551:     } else {
552:       my_coor = coor;
553:     }
554:     PetscCall(PetscMPIIntCast(bs * my_npoints, &imy));
555:     PetscCallMPI(MPI_Bcast(my_coor, imy, MPIU_REAL, 0, comm));
556:   } else {
557:     my_npoints = npoints;
558:     my_coor    = coor;
559:   }

561:   /* determine the number of points living in the bounding box */
562:   n_estimate = 0;
563:   for (i = 0; i < my_npoints; i++) {
564:     PetscBool point_inside = PETSC_TRUE;

566:     for (b = 0; b < bs; b++) {
567:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
568:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
569:     }
570:     if (point_inside) n_estimate++;
571:   }

573:   /* create candidate list */
574:   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
575:   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
576:   PetscCall(VecSetBlockSize(pos, bs));
577:   PetscCall(VecSetFromOptions(pos));
578:   PetscCall(VecGetArray(pos, &_pos));

580:   n_estimate = 0;
581:   for (i = 0; i < my_npoints; i++) {
582:     PetscBool point_inside = PETSC_TRUE;

584:     for (b = 0; b < bs; b++) {
585:       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
586:       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
587:     }
588:     if (point_inside) {
589:       for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = my_coor[bs * i + b];
590:       n_estimate++;
591:     }
592:   }
593:   PetscCall(VecRestoreArray(pos, &_pos));

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

598:   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
599:   n_found = 0;
600:   for (p = 0; p < n_estimate; p++) {
601:     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
602:   }

604:   /* adjust size */
605:   if (mode == ADD_VALUES) {
606:     PetscCall(DMSwarmGetLocalSize(sw, &n_curr));
607:     n_new_est = n_curr + n_found;
608:     PetscCall(DMSwarmSetLocalSizes(sw, n_new_est, -1));
609:   }
610:   if (mode == INSERT_VALUES) {
611:     n_curr    = 0;
612:     n_new_est = n_found;
613:     PetscCall(DMSwarmSetLocalSizes(sw, n_new_est, -1));
614:   }

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

633:   if (redundant) {
634:     if (rank > 0) PetscCall(PetscFree(my_coor));
635:   }
636:   PetscCall(PetscSFDestroy(&sfcell));
637:   PetscCall(VecDestroy(&pos));
638:   PetscFunctionReturn(PETSC_SUCCESS);
639: }

641: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
642: extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);

644: /*@
645:   DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell

647:   Not Collective

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

654:   Level: beginner

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

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

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

668: .seealso: `DMSWARM`, `DMSwarmPICLayoutType`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
669: @*/
670: PetscErrorCode DMSwarmInsertPointsUsingCellDM(DM dm, DMSwarmPICLayoutType layout_type, PetscInt fill_param)
671: {
672:   DM        celldm;
673:   PetscBool isDA, isPLEX;

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

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

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

693:   Not Collective

695:   Input Parameters:
696: + dm      - the `DMSWARM`
697: . npoints - the number of points to insert in each cell
698: - xi      - the coordinates (defined in the local coordinate system for each cell) to insert

700:   Level: beginner

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

715: .seealso: `DMSWARM`, `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
716: @*/
717: PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
718: {
719:   DM        celldm;
720:   PetscBool isDA, isPLEX;

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

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

736:   Not Collective

738:   Input Parameter:
739: . sw - the `DMSWARM`

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

745:   Level: beginner

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

750: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
751: @*/
752: PetscErrorCode DMSwarmCreatePointPerCellCount(DM sw, PetscInt *ncells, PetscInt **count)
753: {
754:   DMSwarmCellDM celldm;
755:   PetscBool     isvalid;
756:   PetscInt      nel;
757:   PetscInt     *sum;
758:   const char   *cellid;

760:   PetscFunctionBegin;
761:   PetscCall(DMSwarmSortGetIsValid(sw, &isvalid));
762:   nel = 0;
763:   if (isvalid) {
764:     PetscInt e;

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

768:     PetscCall(PetscMalloc1(nel, &sum));
769:     for (e = 0; e < nel; e++) PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, e, &sum[e]));
770:   } else {
771:     DM        dm;
772:     PetscBool isda, isplex, isshell;
773:     PetscInt  p, npoints;
774:     PetscInt *swarm_cellid;

776:     /* get the number of cells */
777:     PetscCall(DMSwarmGetCellDMActive(sw, &celldm));
778:     PetscCall(DMSwarmCellDMGetDM(celldm, &dm));
779:     PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMDA, &isda));
780:     PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMPLEX, &isplex));
781:     PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMSHELL, &isshell));
782:     if (isda) {
783:       PetscInt        _nel, _npe;
784:       const PetscInt *_element;

786:       PetscCall(DMDAGetElements(dm, &_nel, &_npe, &_element));
787:       nel = _nel;
788:       PetscCall(DMDARestoreElements(dm, &_nel, &_npe, &_element));
789:     } else if (isplex) {
790:       PetscInt ps, pe;

792:       PetscCall(DMPlexGetHeightStratum(dm, 0, &ps, &pe));
793:       nel = pe - ps;
794:     } else if (isshell) {
795:       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);

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

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

819: /*@
820:   DMSwarmGetNumSpecies - Get the number of particle species

822:   Not Collective

824:   Input Parameter:
825: . sw - the `DMSWARM`

827:   Output Parameters:
828: . Ns - the number of species

830:   Level: intermediate

832: .seealso: `DMSWARM`, `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
833: @*/
834: PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
835: {
836:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

838:   PetscFunctionBegin;
839:   *Ns = swarm->Ns;
840:   PetscFunctionReturn(PETSC_SUCCESS);
841: }

843: /*@
844:   DMSwarmSetNumSpecies - Set the number of particle species

846:   Not Collective

848:   Input Parameters:
849: + sw - the `DMSWARM`
850: - Ns - the number of species

852:   Level: intermediate

854: .seealso: `DMSWARM`, `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
855: @*/
856: PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
857: {
858:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

860:   PetscFunctionBegin;
861:   swarm->Ns = Ns;
862:   PetscFunctionReturn(PETSC_SUCCESS);
863: }

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

868:   Not Collective

870:   Input Parameter:
871: . sw - the `DMSWARM`

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

876:   Level: intermediate

878: .seealso: `DMSWARM`, `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`, `PetscSimplePointFn`
879: @*/
880: PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFn **coordFunc)
881: {
882:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

884:   PetscFunctionBegin;
886:   PetscAssertPointer(coordFunc, 2);
887:   *coordFunc = swarm->coordFunc;
888:   PetscFunctionReturn(PETSC_SUCCESS);
889: }

891: /*@C
892:   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions

894:   Not Collective

896:   Input Parameters:
897: + sw        - the `DMSWARM`
898: - coordFunc - the function setting initial particle positions, see `PetscSimplePointFn` for the calling sequence

900:   Level: intermediate

902: .seealso: `DMSWARM`, `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`, `PetscSimplePointFn`
903: @*/
904: PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFn *coordFunc)
905: {
906:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

908:   PetscFunctionBegin;
911:   swarm->coordFunc = coordFunc;
912:   PetscFunctionReturn(PETSC_SUCCESS);
913: }

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

918:   Not Collective

920:   Input Parameter:
921: . sw - the `DMSWARM`

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

926:   Level: intermediate

928: .seealso: `DMSWARM`, `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`, `PetscSimplePointFn`
929: @*/
930: PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFn **velFunc)
931: {
932:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

934:   PetscFunctionBegin;
936:   PetscAssertPointer(velFunc, 2);
937:   *velFunc = swarm->velFunc;
938:   PetscFunctionReturn(PETSC_SUCCESS);
939: }

941: /*@C
942:   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities

944:   Not Collective

946:   Input Parameters:
947: + sw      - the `DMSWARM`
948: - velFunc - the function setting initial particle velocities, see `PetscSimplePointFn` for the calling sequence

950:   Level: intermediate

952: .seealso: `DMSWARM`, `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`, `PetscSimplePointFn`
953: @*/
954: PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFn *velFunc)
955: {
956:   DM_Swarm *swarm = (DM_Swarm *)sw->data;

958:   PetscFunctionBegin;
961:   swarm->velFunc = velFunc;
962:   PetscFunctionReturn(PETSC_SUCCESS);
963: }

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

968:   Not Collective

970:   Input Parameters:
971: + sw      - The `DMSWARM`
972: . N       - The target number of particles
973: - density - The density field for the particle layout, normalized to unity

975:   Level: advanced

977:   Note:
978:   One particle will be created for each species.

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

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

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

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

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

1050:   Not Collective

1052:   Input Parameter:
1053: . sw - The `DMSWARM`

1055:   Level: advanced

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

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

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

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

1102:   Not Collective

1104:   Input Parameter:
1105: . sw - The `DMSWARM`

1107:   Level: advanced

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

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

1127:   PetscFunctionBeginUser;
1128:   PetscCall(DMGetDimension(sw, &dim));
1129:   PetscCall(DMSwarmGetLocalSize(sw, &Np));
1130:   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
1131:   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));

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

1137:   PetscCall(DMSwarmGetField(sw, coordFields[0], &bs, &dtype, (void **)&x));
1138:   PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight));
1139:   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
1140:   if (coordFunc) {
1141:     PetscCall(DMGetApplicationContext(sw, &ctx));
1142:     for (p = 0; p < Np; ++p) {
1143:       PetscScalar X[3];

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

1156:     PetscCall(DMSwarmGetCellDM(sw, &dm));
1157:     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
1158:     PetscCall(DMGetApplicationContext(sw, &ctx));

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

1170:       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
1171:       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ));
1172:       for (q = 0; q < Npc; ++q) {
1173:         const PetscInt p = pidx[q];
1174:         PetscReal      xref[3];

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

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

1191:   PetscCall(DMSwarmMigrate(sw, removePoints));
1192:   PetscCall(DMLocalizeCoordinates(sw));
1193:   PetscFunctionReturn(PETSC_SUCCESS);
1194: }

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

1199:   Collective

1201:   Input Parameters:
1202: + sw      - The `DMSWARM` object
1203: . sampler - A function which uniformly samples the velocity PDF
1204: - v0      - The velocity scale for nondimensionalization for each species

1206:   Level: advanced

1208:   Note:
1209:   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.

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

1221:   PetscFunctionBegin;
1222:   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));

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

1236:       PetscCall((*velFunc)(dim, 0., NULL, p, vel, ctx));
1237:       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
1238:     }
1239:   } else {
1240:     PetscRandom rnd;

1242:     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
1243:     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
1244:     PetscCall(PetscRandomSetFromOptions(rnd));

1246:     for (p = 0; p < Np; ++p) {
1247:       PetscInt  s = species[p], d;
1248:       PetscReal a[3], vel[3];

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

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

1264:   Collective

1266:   Input Parameters:
1267: + sw - The `DMSWARM` object
1268: - v0 - The velocity scale for nondimensionalization for each species

1270:   Level: advanced

1272: .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
1273: @*/
1274: PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
1275: {
1276:   PetscProbFn *sampler;
1277:   PetscInt     dim;
1278:   const char  *prefix;
1279:   char         funcname[PETSC_MAX_PATH_LEN];
1280:   PetscBool    flg;

1282:   PetscFunctionBegin;
1283:   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
1284:   PetscCall(PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg));
1285:   PetscOptionsEnd();
1286:   if (flg) {
1287:     PetscSimplePointFn *velFunc;

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

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

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

1330:   fegeom.dim      = dim;
1331:   fegeom.dimEmbed = dE;
1332:   fegeom.v        = v0;
1333:   fegeom.xi       = v0ref;
1334:   fegeom.J        = J;
1335:   fegeom.invJ     = invJ;
1336:   fegeom.detJ     = &detJ;

1338:   PetscCall(DMSwarmGetLocalSize(dm, &Np));
1339:   PetscCall(VecGetLocalSize(X, &n));
1340:   PetscCheck(n == Np, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Output vector local size %" PetscInt_FMT " != %" PetscInt_FMT " number of local particles", n, Np);
1341:   PetscCall(DMSwarmGetCellDMActive(dm, &celldm));
1342:   PetscCall(DMSwarmCellDMGetCoordinateFields(celldm, &Nfc, &coordFields));
1343:   PetscCall(DMSwarmCellDMGetFields(celldm, &Nf, &fields));

1345:   PetscCall(PetscMalloc2(Nfc, &coordVals, Nfc, &cbs));
1346:   for (PetscInt i = 0; i < Nfc; ++i) PetscCall(DMSwarmGetField(dm, coordFields[i], &cbs[i], NULL, (void **)&coordVals[i]));
1347:   PetscCall(PetscMalloc2(Nf, &vals, Nfc, &bs));
1348:   for (PetscInt i = 0; i < Nf; ++i) {
1349:     PetscCall(DMSwarmGetField(dm, fields[i], &bs[i], NULL, (void **)&vals[i]));
1350:     totbs += bs[i];
1351:   }

1353:   PetscCall(DMSwarmSortGetAccess(dm));
1354:   for (PetscInt cell = cStart; cell < cEnd; ++cell) {
1355:     PetscInt *pindices, Npc;

1357:     PetscCall(DMSwarmSortGetPointsPerCell(dm, cell, &Npc, &pindices));
1358:     maxC = PetscMax(maxC, Npc);
1359:     PetscCall(DMSwarmSortRestorePointsPerCell(dm, cell, &Npc, &pindices));
1360:   }
1361:   PetscCall(PetscMalloc3(maxC * dim, &xi, maxC * totbs, &val, Nfu, &T));
1362:   PetscCall(VecGetArray(X, &a));
1363:   {
1364:     for (PetscInt cell = cStart; cell < cEnd; ++cell) {
1365:       PetscInt *pindices, Npc;

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

1371:       PetscScalar *closure = NULL;
1372:       PetscInt     Ncl;

1374:       // Get fields from input vector and auxiliary fields from swarm
1375:       for (PetscInt p = 0; p < Npc; ++p) {
1376:         PetscReal xr[8];
1377:         PetscInt  off;

1379:         off = 0;
1380:         for (PetscInt i = 0; i < Nfc; ++i) {
1381:           for (PetscInt b = 0; b < cbs[i]; ++b, ++off) xr[off] = coordVals[i][pindices[p] * cbs[i] + b];
1382:         }
1383:         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);
1384:         CoordinatesRealToRef(dE, dim, fegeom.xi, fegeom.v, fegeom.invJ, xr, &xi[p * dim]);
1385:         off = 0;
1386:         for (PetscInt i = 0; i < Nf; ++i) {
1387:           for (PetscInt b = 0; b < bs[i]; ++b, ++off) val[p * totbs + off] = vals[i][pindices[p] * bs[i] + b];
1388:         }
1389:         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);
1390:       }
1391:       PetscCall(DMPlexVecGetClosure(dmIn, NULL, U, cell, &Ncl, &closure));
1392:       for (PetscInt field = 0; field < Nfu; ++field) {
1393:         PetscFE fe;

1395:         PetscCall(PetscDSGetDiscretization(ds, field, (PetscObject *)&fe));
1396:         PetscCall(PetscFECreateTabulation(fe, 1, Npc, xi, 1, &T[field]));
1397:       }
1398:       for (PetscInt p = 0; p < Npc; ++p) {
1399:         // Get fields from input vector
1400:         PetscCall(PetscFEEvaluateFieldJets_Internal(ds, Nfu, 0, p, T, &fegeom, closure, NULL, u, u_x, NULL));
1401:         (*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]]);
1402:       }
1403:       PetscCall(DMSwarmSortRestorePointsPerCell(dm, cell, &Npc, &pindices));
1404:       PetscCall(DMPlexVecRestoreClosure(dmIn, NULL, U, cell, &Ncl, &closure));
1405:       for (PetscInt field = 0; field < Nfu; ++field) PetscCall(PetscTabulationDestroy(&T[field]));
1406:     }
1407:   }
1408:   for (PetscInt i = 0; i < Nfc; ++i) PetscCall(DMSwarmRestoreField(dm, coordFields[i], &cbs[i], NULL, (void **)&coordVals[i]));
1409:   for (PetscInt i = 0; i < Nf; ++i) PetscCall(DMSwarmRestoreField(dm, fields[i], &bs[i], NULL, (void **)&vals[i]));
1410:   PetscCall(VecRestoreArray(X, &a));
1411:   PetscCall(DMSwarmSortRestoreAccess(dm));
1412:   PetscCall(PetscFree3(xi, val, T));
1413:   PetscCall(PetscFree3(v0, J, invJ));
1414:   PetscCall(PetscFree2(coordVals, cbs));
1415:   PetscCall(PetscFree2(vals, bs));
1416:   PetscFunctionReturn(PETSC_SUCCESS);
1417: }