1: #include "petscdm.h"
  2: #include "petscsystypes.h"
  3: #include <petscdmda.h>
  4: #include <petscdmplex.h>
  5: #include <petsc/private/dmswarmimpl.h>

  7: static int sort_CompareSwarmPoint(const void *dataA, const void *dataB)
  8: {
  9:   SwarmPoint *pointA = (SwarmPoint *)dataA;
 10:   SwarmPoint *pointB = (SwarmPoint *)dataB;

 12:   if (pointA->cell_index < pointB->cell_index) {
 13:     return -1;
 14:   } else if (pointA->cell_index > pointB->cell_index) {
 15:     return 1;
 16:   } else {
 17:     return 0;
 18:   }
 19: }

 21: static PetscErrorCode DMSwarmSortApplyCellIndexSort(DMSwarmSort ctx)
 22: {
 23:   PetscFunctionBegin;
 24:   if (ctx->list) qsort(ctx->list, ctx->npoints, sizeof(SwarmPoint), sort_CompareSwarmPoint);
 25:   PetscFunctionReturn(PETSC_SUCCESS);
 26: }

 28: static PetscErrorCode DMSwarmSortCreate(DMSwarmSort *_ctx)
 29: {
 30:   DMSwarmSort ctx;

 32:   PetscFunctionBegin;
 33:   PetscCall(PetscNew(&ctx));
 34:   ctx->isvalid = PETSC_FALSE;
 35:   ctx->ncells  = 0;
 36:   ctx->npoints = 0;
 37:   PetscCall(PetscMalloc1(1, &ctx->pcell_offsets));
 38:   PetscCall(PetscMalloc1(1, &ctx->list));
 39:   *_ctx = ctx;
 40:   PetscFunctionReturn(PETSC_SUCCESS);
 41: }

 43: static PetscErrorCode DMSwarmSortSetup(DMSwarmSort ctx, DM dm, PetscInt ncells)
 44: {
 45:   PetscInt *swarm_cellid;
 46:   PetscInt  p, npoints;
 47:   PetscInt  tmp, c, count;

 49:   PetscFunctionBegin;
 50:   if (!ctx) PetscFunctionReturn(PETSC_SUCCESS);
 51:   if (ctx->isvalid) PetscFunctionReturn(PETSC_SUCCESS);

 53:   PetscCall(PetscLogEventBegin(DMSWARM_Sort, 0, 0, 0, 0));
 54:   /* check the number of cells */
 55:   if (ncells != ctx->ncells) {
 56:     PetscCall(PetscRealloc(sizeof(PetscInt) * (ncells + 1), &ctx->pcell_offsets));
 57:     ctx->ncells = ncells;
 58:   }
 59:   PetscCall(PetscArrayzero(ctx->pcell_offsets, ctx->ncells + 1));

 61:   /* get the number of points */
 62:   PetscCall(DMSwarmGetLocalSize(dm, &npoints));
 63:   if (npoints != ctx->npoints) {
 64:     PetscCall(PetscRealloc(sizeof(SwarmPoint) * npoints, &ctx->list));
 65:     ctx->npoints = npoints;
 66:   }
 67:   PetscCall(PetscArrayzero(ctx->list, npoints));

 69:   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
 70:   for (p = 0; p < ctx->npoints; p++) {
 71:     ctx->list[p].point_index = p;
 72:     ctx->list[p].cell_index  = swarm_cellid[p];
 73:   }
 74:   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
 75:   PetscCall(DMSwarmSortApplyCellIndexSort(ctx));

 77:   /* sum points per cell */
 78:   for (p = 0; p < ctx->npoints; p++) ctx->pcell_offsets[ctx->list[p].cell_index]++;

 80:   /* create offset list */
 81:   count = 0;
 82:   for (c = 0; c < ctx->ncells; c++) {
 83:     tmp                   = ctx->pcell_offsets[c];
 84:     ctx->pcell_offsets[c] = count;
 85:     count                 = count + tmp;
 86:   }
 87:   ctx->pcell_offsets[c] = count;

 89:   ctx->isvalid = PETSC_TRUE;
 90:   PetscCall(PetscLogEventEnd(DMSWARM_Sort, 0, 0, 0, 0));
 91:   PetscFunctionReturn(PETSC_SUCCESS);
 92: }

 94: PetscErrorCode DMSwarmSortDestroy(DMSwarmSort *_ctx)
 95: {
 96:   DMSwarmSort ctx;

 98:   PetscFunctionBegin;
 99:   if (!_ctx) PetscFunctionReturn(PETSC_SUCCESS);
100:   if (!*_ctx) PetscFunctionReturn(PETSC_SUCCESS);
101:   ctx = *_ctx;
102:   if (ctx->list) PetscCall(PetscFree(ctx->list));
103:   if (ctx->pcell_offsets) PetscCall(PetscFree(ctx->pcell_offsets));
104:   PetscCall(PetscFree(ctx));
105:   *_ctx = NULL;
106:   PetscFunctionReturn(PETSC_SUCCESS);
107: }

109: /*@
110:   DMSwarmSortGetNumberOfPointsPerCell - Returns the number of points in a cell

112:   Not Collective

114:   Input Parameters:
115: + dm - a `DMSWARM` objects
116: - e  - the index of the cell

118:   Output Parameter:
119: . npoints - the number of points in the cell

121:   Level: advanced

123:   Notes:
124:   You must call `DMSwarmSortGetAccess()` before you can call `DMSwarmSortGetNumberOfPointsPerCell()`

126: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSortGetAccess()`, `DMSwarmSortGetPointsPerCell()`
127: @*/
128: PetscErrorCode DMSwarmSortGetNumberOfPointsPerCell(DM dm, PetscInt e, PetscInt *npoints)
129: {
130:   DM_Swarm   *swarm = (DM_Swarm *)dm->data;
131:   PetscInt    points_per_cell;
132:   DMSwarmSort ctx;

134:   PetscFunctionBegin;
135:   ctx = swarm->sort_context;
136:   PetscCheck(ctx, PetscObjectComm((PetscObject)dm), PETSC_ERR_USER, "The DMSwarmSort context has not been created. Must call DMSwarmSortGetAccess() first");
137:   PetscCheck(ctx->isvalid, PETSC_COMM_SELF, PETSC_ERR_USER, "SwarmPointSort container is not valid. Must call DMSwarmSortGetAccess() first");
138:   PetscCheck(e < ctx->ncells, PETSC_COMM_SELF, PETSC_ERR_USER, "Cell index (%" PetscInt_FMT ") is greater than max number of local cells (%" PetscInt_FMT ")", e, ctx->ncells);
139:   PetscCheck(e >= 0, PETSC_COMM_SELF, PETSC_ERR_USER, "Cell index (%" PetscInt_FMT ") cannot be negative", e);
140:   points_per_cell = ctx->pcell_offsets[e + 1] - ctx->pcell_offsets[e];
141:   *npoints        = points_per_cell;
142:   PetscFunctionReturn(PETSC_SUCCESS);
143: }

145: /*@C
146:   DMSwarmSortGetPointsPerCell - Creates an array of point indices for all points in a cell

148:   Not Collective

150:   Input Parameters:
151: + dm      - a `DMSWARM` object
152: . e       - the index of the cell
153: . npoints - the number of points in the cell
154: - pidlist - array of the indices identifying all points in cell e

156:   Level: advanced

158:   Note:
159:   You must call `DMSwarmSortGetAccess()` before you can call `DMSwarmSortGetPointsPerCell()`, and call `DMSwarmRestorePointsPerCell()` afterwards

161: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmRestorePointsPerCell()`, `DMSwarmSortGetAccess()`, `DMSwarmSortGetNumberOfPointsPerCell()`
162: @*/
163: PetscErrorCode DMSwarmSortGetPointsPerCell(DM dm, PetscInt e, PetscInt *npoints, PetscInt **pidlist)
164: {
165:   DM_Swarm   *swarm = (DM_Swarm *)dm->data;
166:   PetscInt    p, pid, pid_unsorted;
167:   DMSwarmSort ctx;

169:   PetscFunctionBegin;
170:   ctx = swarm->sort_context;
171:   PetscCheck(ctx, PetscObjectComm((PetscObject)dm), PETSC_ERR_USER, "The DMSwarmSort context has not been created. Must call DMSwarmSortGetAccess() first");
172:   PetscCall(DMSwarmSortGetNumberOfPointsPerCell(dm, e, npoints));
173:   PetscCall(DMGetWorkArray(dm, *npoints, MPIU_SCALAR, pidlist));
174:   for (p = 0; p < *npoints; p++) {
175:     pid           = ctx->pcell_offsets[e] + p;
176:     pid_unsorted  = ctx->list[pid].point_index;
177:     (*pidlist)[p] = pid_unsorted;
178:   }
179:   PetscFunctionReturn(PETSC_SUCCESS);
180: }

182: /*@C
183:   DMSwarmSortRestorePointsPerCell - Restores an array of point indices for all points in a cell

185:   Not Collective

187:   Input Parameters:
188: + dm      - a `DMSWARM` object
189: . e       - the index of the cell
190: . npoints - the number of points in the cell
191: - pidlist - array of the indices identifying all points in cell e

193:   Level: advanced

195:   Note:
196:   You must call `DMSwarmSortGetAccess()` and `DMSwarmSortGetPointsPerCell()` before you can call `DMSwarmSortRestorePointsPerCell()`

198: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSortGetPointsPerCell()`, `DMSwarmSortGetAccess()`, `DMSwarmSortGetNumberOfPointsPerCell()`
199: @*/
200: PetscErrorCode DMSwarmSortRestorePointsPerCell(DM dm, PetscInt e, PetscInt *npoints, PetscInt **pidlist)
201: {
202:   PetscFunctionBegin;
203:   PetscCall(DMRestoreWorkArray(dm, *npoints, MPIU_SCALAR, pidlist));
204:   PetscFunctionReturn(PETSC_SUCCESS);
205: }

207: /*@
208:   DMSwarmSortGetAccess - Setups up a `DMSWARM` point sort context for efficient traversal of points within a cell

210:   Not Collective

212:   Input Parameter:
213: . dm - a `DMSWARM` object

215:   Level: advanced

217:   Notes:
218:   Calling `DMSwarmSortGetAccess()` creates a list which enables easy identification of all points contained in a
219:   given cell. This method does not explicitly sort the data within the `DMSWARM` based on the cell index associated
220:   with a `DMSWARM` point.

222:   The sort context is valid only for the `DMSWARM` points defined at the time when `DMSwarmSortGetAccess()` was called.
223:   For example, suppose the swarm contained NP points when `DMSwarmSortGetAccess()` was called. If the user subsequently
224:   adds 10 additional points to the swarm, the sort context is still valid, but only for the first NP points.
225:   The indices associated with the 10 new additional points will not be contained within the sort context.
226:   This means that the user can still safely perform queries via `DMSwarmSortGetPointsPerCell()` and
227:   `DMSwarmSortGetPointsPerCell()`, however the results return will be based on the first NP points.

229:   If any` DMSWARM` re-sizing method is called after `DMSwarmSortGetAccess()` which modifies any of the first NP entries
230:   in the `DMSWARM`, the sort context will become invalid. Currently there are no guards to prevent the user from
231:   invalidating the sort context. For this reason, we highly recommend you do not use `DMSwarmRemovePointAtIndex()` in
232:   between calls to `DMSwarmSortGetAccess()` and `DMSwarmSortRestoreAccess()`.

234:   To facilitate safe removal of points using the sort context, we suggest a "two pass" strategy in which the
235:   first pass "marks" points for removal, and the second pass actually removes the points from the `DMSWARM`

237:   You must call `DMSwarmSortGetAccess()` before you can call `DMSwarmSortGetPointsPerCell()` or `DMSwarmSortGetNumberOfPointsPerCell()`

239:   The sort context may become invalid if any re-sizing methods are applied which alter the first NP points
240:   within swarm at the time `DMSwarmSortGetAccess()` was called.

242:   You must call `DMSwarmSortRestoreAccess()` when you no longer need access to the sort context

244: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSortRestoreAccess()`
245: @*/
246: PetscErrorCode DMSwarmSortGetAccess(DM dm)
247: {
248:   DM_Swarm *swarm = (DM_Swarm *)dm->data;
249:   PetscInt  ncells;
250:   DM        celldm;
251:   PetscBool isda, isplex, isshell;

253:   PetscFunctionBegin;
254:   if (!swarm->sort_context) PetscCall(DMSwarmSortCreate(&swarm->sort_context));

256:   /* get the number of cells */
257:   PetscCall(DMSwarmGetCellDM(dm, &celldm));
258:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isda));
259:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isplex));
260:   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMSHELL, &isshell));
261:   ncells = 0;
262:   if (isda) {
263:     PetscInt        nel, npe;
264:     const PetscInt *element;

266:     PetscCall(DMDAGetElements(celldm, &nel, &npe, &element));
267:     ncells = nel;
268:     PetscCall(DMDARestoreElements(celldm, &nel, &npe, &element));
269:   } else if (isplex) {
270:     PetscInt ps, pe;

272:     PetscCall(DMPlexGetHeightStratum(celldm, 0, &ps, &pe));
273:     ncells = pe - ps;
274:   } else if (isshell) {
275:     PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);

277:     PetscCall(PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
278:     if (method_DMShellGetNumberOfCells) {
279:       PetscCall(method_DMShellGetNumberOfCells(celldm, &ncells));
280:     } else
281:       SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for the DMSHELL object. User must provide a method via PetscObjectComposeFunction( (PetscObject)shelldm, \"DMGetNumberOfCells_C\", your_function_to_compute_number_of_cells);");
282:   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");

284:   /* setup */
285:   PetscCall(DMSwarmSortSetup(swarm->sort_context, dm, ncells));
286:   PetscFunctionReturn(PETSC_SUCCESS);
287: }

289: /*@
290:   DMSwarmSortRestoreAccess - Invalidates the `DMSWARM` point sorting context previously computed with `DMSwarmSortGetAccess()`

292:   Not Collective

294:   Input Parameter:
295: . dm - a `DMSWARM` object

297:   Level: advanced

299:   Note:
300:   You must call `DMSwarmSortGetAccess()` before calling `DMSwarmSortRestoreAccess()`

302: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSortGetAccess()`
303: @*/
304: PetscErrorCode DMSwarmSortRestoreAccess(DM dm)
305: {
306:   DM_Swarm *swarm = (DM_Swarm *)dm->data;

308:   PetscFunctionBegin;
309:   if (!swarm->sort_context) PetscFunctionReturn(PETSC_SUCCESS);
310:   PetscCheck(swarm->sort_context->isvalid, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "You must call DMSwarmSortGetAccess() before calling DMSwarmSortRestoreAccess()");
311:   swarm->sort_context->isvalid = PETSC_FALSE;
312:   PetscFunctionReturn(PETSC_SUCCESS);
313: }

315: /*@
316:   DMSwarmSortGetIsValid - Gets the isvalid flag associated with a `DMSWARM` point sorting context

318:   Not Collective

320:   Input Parameter:
321: . dm - a `DMSWARM` object

323:   Output Parameter:
324: . isvalid - flag indicating whether the sort context is up-to-date

326:   Level: advanced

328: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSortGetAccess()`
329: @*/
330: PetscErrorCode DMSwarmSortGetIsValid(DM dm, PetscBool *isvalid)
331: {
332:   DM_Swarm *swarm = (DM_Swarm *)dm->data;

334:   PetscFunctionBegin;
335:   if (!swarm->sort_context) {
336:     *isvalid = PETSC_FALSE;
337:     PetscFunctionReturn(PETSC_SUCCESS);
338:   }
339:   *isvalid = swarm->sort_context->isvalid;
340:   PetscFunctionReturn(PETSC_SUCCESS);
341: }

343: /*@
344:   DMSwarmSortGetSizes - Gets the sizes associated with a `DMSWARM` point sorting context

346:   Not Collective

348:   Input Parameter:
349: . dm - a `DMSWARM` object

351:   Output Parameters:
352: + ncells  - number of cells within the sort context (pass `NULL` to ignore)
353: - npoints - number of points used to create the sort context (pass `NULL` to ignore)

355:   Level: advanced

357: .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSortGetAccess()`
358: @*/
359: PetscErrorCode DMSwarmSortGetSizes(DM dm, PetscInt *ncells, PetscInt *npoints)
360: {
361:   DM_Swarm *swarm = (DM_Swarm *)dm->data;

363:   PetscFunctionBegin;
364:   if (!swarm->sort_context) {
365:     if (ncells) *ncells = 0;
366:     if (npoints) *npoints = 0;
367:     PetscFunctionReturn(PETSC_SUCCESS);
368:   }
369:   if (ncells) *ncells = swarm->sort_context->ncells;
370:   if (npoints) *npoints = swarm->sort_context->npoints;
371:   PetscFunctionReturn(PETSC_SUCCESS);
372: }