Actual source code: swarmpic.c

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

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

 11: PetscClassId DMSWARMCELLDM_CLASSID;

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

 16:   Collective

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

 21:   Level: advanced

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

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

 49:   Collective

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

 55:   Level: advanced

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

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

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

 90:   Not Collective

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

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

 98:   Level: intermediate

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

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

114:   Not Collective

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

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

123:   Level: intermediate

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

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

145:   Not Collective

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

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

154:   Level: intermediate

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

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

176:   Not Collective

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

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

184:   Level: intermediate

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

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

200:   Not Collective

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

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

208:   Level: intermediate

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

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

224:   Not Collective

226:   Input Parameters:
227: + celldm - The `DMSwarmCellDM` object
228: - sort   - The `DMSwarmSort` object

230:   Level: intermediate

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

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

246:   Not Collective

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

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

255:   Level: intermediate

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

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

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

278:   Collective

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

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

290:   Level: advanced

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

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

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

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

327:   Collective

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

336:   Level: beginner

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

481:   Collective

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

490:   Level: beginner

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

648:   Not Collective

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

655:   Level: beginner

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

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

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

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

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

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

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

694:   Not Collective

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

701:   Level: beginner

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

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

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

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

737:   Not Collective

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

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

746:   Level: beginner

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

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

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

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

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

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

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

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

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

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

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

823:   Not Collective

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

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

831:   Level: intermediate

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

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

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

847:   Not Collective

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

853:   Level: intermediate

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

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

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

869:   Not Collective

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

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

877:   Level: intermediate

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

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

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

895:   Not Collective

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

901:   Level: intermediate

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

909:   PetscFunctionBegin;
912:   swarm->coordFunc = coordFunc;
913:   PetscFunctionReturn(PETSC_SUCCESS);
914: }

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

919:   Not Collective

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

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

927:   Level: intermediate

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

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

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

945:   Not Collective

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

951:   Level: intermediate

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

959:   PetscFunctionBegin;
962:   swarm->velFunc = velFunc;
963:   PetscFunctionReturn(PETSC_SUCCESS);
964: }

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

969:   Not Collective

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

976:   Level: advanced

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

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

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

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

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

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

1051:   Not Collective

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

1056:   Level: advanced

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

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

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

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

1103:   Not Collective

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

1108:   Level: advanced

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

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

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

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

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

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

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

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

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

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

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

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

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

1200:   Collective

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

1207:   Level: advanced

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

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

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

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

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

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

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

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

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

1265:   Collective

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

1271:   Level: advanced

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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