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: }