xref: /petsc/src/dm/impls/swarm/swarmpic.c (revision c77c71ff2d9eaa2c74538bf9bf94eff01b512dbf)
1 #define PETSCDM_DLL
2 #include <petsc/private/dmswarmimpl.h> /*I   "petscdmswarm.h"   I*/
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"
8 
9 #include <petsc/private/petscfeimpl.h> /* For CoordinatesRefToReal() */
10 
11 /*
12  Error checking to ensure the swarm type is correct and that a cell DM has been set
13 */
14 #define DMSWARMPICVALID(dm) \
15   do { \
16     DM_Swarm *_swarm = (DM_Swarm *)(dm)->data; \
17     PetscCheck(_swarm->swarm_type == DMSWARM_PIC, PetscObjectComm((PetscObject)(dm)), PETSC_ERR_SUP, "Valid only for DMSwarm-PIC. You must call DMSwarmSetType(dm,DMSWARM_PIC)"); \
18     PetscCheck(_swarm->dmcell, PetscObjectComm((PetscObject)(dm)), PETSC_ERR_SUP, "Valid only for DMSwarmPIC if the cell DM is set. You must call DMSwarmSetCellDM(dm,celldm)"); \
19   } while (0)
20 
21 /* Coordinate insertition/addition API */
22 /*@C
23    DMSwarmSetPointsUniformCoordinates - Set point coordinates in a `DMSWARM` on a regular (ijk) grid
24 
25    Collective
26 
27    Input parameters:
28 +  dm - the `DMSWARM`
29 .  min - minimum coordinate values in the x, y, z directions (array of length dim)
30 .  max - maximum coordinate values in the x, y, z directions (array of length dim)
31 .  npoints - number of points in each spatial direction (array of length dim)
32 -  mode - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)
33 
34    Level: beginner
35 
36    Notes:
37    When using mode = `INSERT_VALUES`, this method will reset the number of particles in the `DMSWARM`
38    to be npoints[0]*npoints[1] (2D) or npoints[0]*npoints[1]*npoints[2] (3D). When using mode = `ADD_VALUES`,
39    new points will be appended to any already existing in the `DMSWARM`
40 
41 .seealso: `DM`, `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
42 @*/
43 PETSC_EXTERN PetscErrorCode DMSwarmSetPointsUniformCoordinates(DM dm, PetscReal min[], PetscReal max[], PetscInt npoints[], InsertMode mode)
44 {
45   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
46   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
47   PetscInt           i, j, k, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
48   Vec                coorlocal;
49   const PetscScalar *_coor;
50   DM                 celldm;
51   PetscReal          dx[3];
52   PetscInt           _npoints[] = {0, 0, 1};
53   Vec                pos;
54   PetscScalar       *_pos;
55   PetscReal         *swarm_coor;
56   PetscInt          *swarm_cellid;
57   PetscSF            sfcell = NULL;
58   const PetscSFNode *LA_sfcell;
59 
60   PetscFunctionBegin;
61   DMSWARMPICVALID(dm);
62   PetscCall(DMSwarmGetCellDM(dm, &celldm));
63   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
64   PetscCall(VecGetSize(coorlocal, &N));
65   PetscCall(VecGetBlockSize(coorlocal, &bs));
66   N = N / bs;
67   PetscCall(VecGetArrayRead(coorlocal, &_coor));
68   for (i = 0; i < N; i++) {
69     for (b = 0; b < bs; b++) {
70       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
71       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
72     }
73   }
74   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));
75 
76   for (b = 0; b < bs; b++) {
77     if (npoints[b] > 1) {
78       dx[b] = (max[b] - min[b]) / ((PetscReal)(npoints[b] - 1));
79     } else {
80       dx[b] = 0.0;
81     }
82     _npoints[b] = npoints[b];
83   }
84 
85   /* determine number of points living in the bounding box */
86   n_estimate = 0;
87   for (k = 0; k < _npoints[2]; k++) {
88     for (j = 0; j < _npoints[1]; j++) {
89       for (i = 0; i < _npoints[0]; i++) {
90         PetscReal xp[] = {0.0, 0.0, 0.0};
91         PetscInt  ijk[3];
92         PetscBool point_inside = PETSC_TRUE;
93 
94         ijk[0] = i;
95         ijk[1] = j;
96         ijk[2] = k;
97         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
98         for (b = 0; b < bs; b++) {
99           if (xp[b] < gmin[b]) point_inside = PETSC_FALSE;
100           if (xp[b] > gmax[b]) point_inside = PETSC_FALSE;
101         }
102         if (point_inside) n_estimate++;
103       }
104     }
105   }
106 
107   /* create candidate list */
108   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
109   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
110   PetscCall(VecSetBlockSize(pos, bs));
111   PetscCall(VecSetFromOptions(pos));
112   PetscCall(VecGetArray(pos, &_pos));
113 
114   n_estimate = 0;
115   for (k = 0; k < _npoints[2]; k++) {
116     for (j = 0; j < _npoints[1]; j++) {
117       for (i = 0; i < _npoints[0]; i++) {
118         PetscReal xp[] = {0.0, 0.0, 0.0};
119         PetscInt  ijk[3];
120         PetscBool point_inside = PETSC_TRUE;
121 
122         ijk[0] = i;
123         ijk[1] = j;
124         ijk[2] = k;
125         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
126         for (b = 0; b < bs; b++) {
127           if (xp[b] < gmin[b]) point_inside = PETSC_FALSE;
128           if (xp[b] > gmax[b]) point_inside = PETSC_FALSE;
129         }
130         if (point_inside) {
131           for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = xp[b];
132           n_estimate++;
133         }
134       }
135     }
136   }
137   PetscCall(VecRestoreArray(pos, &_pos));
138 
139   /* locate points */
140   PetscCall(DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell));
141   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
142   n_found = 0;
143   for (p = 0; p < n_estimate; p++) {
144     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
145   }
146 
147   /* adjust size */
148   if (mode == ADD_VALUES) {
149     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
150     n_new_est = n_curr + n_found;
151     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
152   }
153   if (mode == INSERT_VALUES) {
154     n_curr    = 0;
155     n_new_est = n_found;
156     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
157   }
158 
159   /* initialize new coords, cell owners, pid */
160   PetscCall(VecGetArrayRead(pos, &_coor));
161   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
162   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
163   n_found = 0;
164   for (p = 0; p < n_estimate; p++) {
165     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
166       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
167       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
168       n_found++;
169     }
170   }
171   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
172   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
173   PetscCall(VecRestoreArrayRead(pos, &_coor));
174 
175   PetscCall(PetscSFDestroy(&sfcell));
176   PetscCall(VecDestroy(&pos));
177   PetscFunctionReturn(PETSC_SUCCESS);
178 }
179 
180 /*@C
181    DMSwarmSetPointCoordinates - Set point coordinates in a `DMSWARM` from a user defined list
182 
183    Collective
184 
185    Input parameters:
186 +  dm - the `DMSWARM`
187 .  npoints - the number of points to insert
188 .  coor - the coordinate values
189 .  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
190 -  mode - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)
191 
192    Level: beginner
193 
194    Notes:
195    If the user has specified `redundant` as `PETSC_FALSE`, the cell `DM` will attempt to locate the coordinates provided by `coor` within
196    its sub-domain. If they any values within `coor` are not located in the sub-domain, they will be ignored and will not get
197    added to the `DMSWARM`.
198 
199 .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`, `DMSwarmSetPointsUniformCoordinates()`
200 @*/
201 PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinates(DM dm, PetscInt npoints, PetscReal coor[], PetscBool redundant, InsertMode mode)
202 {
203   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
204   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
205   PetscInt           i, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
206   Vec                coorlocal;
207   const PetscScalar *_coor;
208   DM                 celldm;
209   Vec                pos;
210   PetscScalar       *_pos;
211   PetscReal         *swarm_coor;
212   PetscInt          *swarm_cellid;
213   PetscSF            sfcell = NULL;
214   const PetscSFNode *LA_sfcell;
215   PetscReal         *my_coor;
216   PetscInt           my_npoints;
217   PetscMPIInt        rank;
218   MPI_Comm           comm;
219 
220   PetscFunctionBegin;
221   DMSWARMPICVALID(dm);
222   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
223   PetscCallMPI(MPI_Comm_rank(comm, &rank));
224 
225   PetscCall(DMSwarmGetCellDM(dm, &celldm));
226   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
227   PetscCall(VecGetSize(coorlocal, &N));
228   PetscCall(VecGetBlockSize(coorlocal, &bs));
229   N = N / bs;
230   PetscCall(VecGetArrayRead(coorlocal, &_coor));
231   for (i = 0; i < N; i++) {
232     for (b = 0; b < bs; b++) {
233       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
234       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
235     }
236   }
237   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));
238 
239   /* broadcast points from rank 0 if requested */
240   if (redundant) {
241     my_npoints = npoints;
242     PetscCallMPI(MPI_Bcast(&my_npoints, 1, MPIU_INT, 0, comm));
243 
244     if (rank > 0) { /* allocate space */
245       PetscCall(PetscMalloc1(bs * my_npoints, &my_coor));
246     } else {
247       my_coor = coor;
248     }
249     PetscCallMPI(MPI_Bcast(my_coor, bs * my_npoints, MPIU_REAL, 0, comm));
250   } else {
251     my_npoints = npoints;
252     my_coor    = coor;
253   }
254 
255   /* determine the number of points living in the bounding box */
256   n_estimate = 0;
257   for (i = 0; i < my_npoints; i++) {
258     PetscBool point_inside = PETSC_TRUE;
259 
260     for (b = 0; b < bs; b++) {
261       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
262       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
263     }
264     if (point_inside) n_estimate++;
265   }
266 
267   /* create candidate list */
268   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
269   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
270   PetscCall(VecSetBlockSize(pos, bs));
271   PetscCall(VecSetFromOptions(pos));
272   PetscCall(VecGetArray(pos, &_pos));
273 
274   n_estimate = 0;
275   for (i = 0; i < my_npoints; i++) {
276     PetscBool point_inside = PETSC_TRUE;
277 
278     for (b = 0; b < bs; b++) {
279       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
280       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
281     }
282     if (point_inside) {
283       for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = my_coor[bs * i + b];
284       n_estimate++;
285     }
286   }
287   PetscCall(VecRestoreArray(pos, &_pos));
288 
289   /* locate points */
290   PetscCall(DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell));
291 
292   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
293   n_found = 0;
294   for (p = 0; p < n_estimate; p++) {
295     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
296   }
297 
298   /* adjust size */
299   if (mode == ADD_VALUES) {
300     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
301     n_new_est = n_curr + n_found;
302     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
303   }
304   if (mode == INSERT_VALUES) {
305     n_curr    = 0;
306     n_new_est = n_found;
307     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
308   }
309 
310   /* initialize new coords, cell owners, pid */
311   PetscCall(VecGetArrayRead(pos, &_coor));
312   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
313   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
314   n_found = 0;
315   for (p = 0; p < n_estimate; p++) {
316     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
317       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
318       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
319       n_found++;
320     }
321   }
322   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
323   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
324   PetscCall(VecRestoreArrayRead(pos, &_coor));
325 
326   if (redundant) {
327     if (rank > 0) PetscCall(PetscFree(my_coor));
328   }
329   PetscCall(PetscSFDestroy(&sfcell));
330   PetscCall(VecDestroy(&pos));
331   PetscFunctionReturn(PETSC_SUCCESS);
332 }
333 
334 extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
335 extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);
336 
337 /*@C
338    DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell
339 
340    Not Collective
341 
342    Input parameters:
343 +  dm - the `DMSWARM`
344 .  layout_type - method used to fill each cell with the cell `DM`
345 -  fill_param - parameter controlling how many points per cell are added (the meaning of this parameter is dependent on the layout type)
346 
347    Level: beginner
348 
349    Notes:
350    The insert method will reset any previous defined points within the `DMSWARM`.
351 
352    When using a `DMDA` both 2D and 3D are supported for all layout types provided you are using `DMDA_ELEMENT_Q1`.
353 
354    When using a `DMPLEX` the following case are supported:
355 .vb
356    (i) DMSWARMPIC_LAYOUT_REGULAR: 2D (triangle),
357    (ii) DMSWARMPIC_LAYOUT_GAUSS: 2D and 3D provided the cell is a tri/tet or a quad/hex,
358    (iii) DMSWARMPIC_LAYOUT_SUBDIVISION: 2D and 3D for quad/hex and 2D tri.
359 .ve
360 
361 .seealso: `DMSWARM`, `DMSwarmPICLayoutType`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
362 @*/
363 PETSC_EXTERN PetscErrorCode DMSwarmInsertPointsUsingCellDM(DM dm, DMSwarmPICLayoutType layout_type, PetscInt fill_param)
364 {
365   DM        celldm;
366   PetscBool isDA, isPLEX;
367 
368   PetscFunctionBegin;
369   DMSWARMPICVALID(dm);
370   PetscCall(DMSwarmGetCellDM(dm, &celldm));
371   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
372   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
373   if (isDA) {
374     PetscCall(private_DMSwarmInsertPointsUsingCellDM_DA(dm, celldm, layout_type, fill_param));
375   } else if (isPLEX) {
376     PetscCall(private_DMSwarmInsertPointsUsingCellDM_PLEX(dm, celldm, layout_type, fill_param));
377   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
378   PetscFunctionReturn(PETSC_SUCCESS);
379 }
380 
381 extern PetscErrorCode private_DMSwarmSetPointCoordinatesCellwise_PLEX(DM, DM, PetscInt, PetscReal *);
382 
383 /*@C
384    DMSwarmSetPointCoordinatesCellwise - Insert point coordinates (defined over the reference cell) within each cell
385 
386    Not Collective
387 
388    Input parameters:
389 +  dm - the `DMSWARM`
390 .  celldm - the cell `DM`
391 .  npoints - the number of points to insert in each cell
392 -  xi - the coordinates (defined in the local coordinate system for each cell) to insert
393 
394  Level: beginner
395 
396  Notes:
397  The method will reset any previous defined points within the `DMSWARM`.
398  Only supported for `DMPLEX`. If you are using a `DMDA` it is recommended to either use
399  `DMSwarmInsertPointsUsingCellDM()`, or extract and set the coordinates yourself the following code
400 .vb
401     PetscReal *coor;
402     DMSwarmGetField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
403     // user code to define the coordinates here
404     DMSwarmRestoreField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
405 .ve
406 
407 .seealso: `DMSWARM`, `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
408 @*/
409 PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
410 {
411   DM        celldm;
412   PetscBool isDA, isPLEX;
413 
414   PetscFunctionBegin;
415   DMSWARMPICVALID(dm);
416   PetscCall(DMSwarmGetCellDM(dm, &celldm));
417   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
418   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
419   PetscCheck(!isDA, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMPLEX. Recommended you use DMSwarmInsertPointsUsingCellDM()");
420   if (isPLEX) {
421     PetscCall(private_DMSwarmSetPointCoordinatesCellwise_PLEX(dm, celldm, npoints, xi));
422   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
423   PetscFunctionReturn(PETSC_SUCCESS);
424 }
425 
426 /* Field projection API */
427 extern PetscErrorCode private_DMSwarmProjectFields_DA(DM swarm, DM celldm, PetscInt project_type, PetscInt nfields, DMSwarmDataField dfield[], Vec vecs[]);
428 extern PetscErrorCode private_DMSwarmProjectFields_PLEX(DM swarm, DM celldm, PetscInt project_type, PetscInt nfields, DMSwarmDataField dfield[], Vec vecs[]);
429 
430 /*@C
431    DMSwarmProjectFields - Project a set of swarm fields onto the cell `DM`
432 
433    Collective
434 
435    Input parameters:
436 +  dm - the `DMSWARM`
437 .  nfields - the number of swarm fields to project
438 .  fieldnames - the textual names of the swarm fields to project
439 .  fields - an array of Vec's of length nfields
440 -  reuse - flag indicating whether the array and contents of fields should be re-used or internally allocated
441 
442    Level: beginner
443 
444    Notes:
445    Currently, the only available projection method consists of
446 .vb
447      phi_i = \sum_{p=0}^{np} N_i(x_p) phi_p dJ / \sum_{p=0}^{np} N_i(x_p) dJ
448    where phi_p is the swarm field at point p,
449      N_i() is the cell DM basis function at vertex i,
450      dJ is the determinant of the cell Jacobian and
451      phi_i is the projected vertex value of the field phi.
452 .ve
453 
454    If `reuse` is `PETSC_FALSE`, this function will allocate the array of `Vec`'s, and each individual `Vec`.
455      The user is responsible for destroying both the array and the individual `Vec` objects.
456 
457    Only swarm fields registered with data type of `PETSC_REAL` can be projected onto the cell `DM`.
458 
459    Only swarm fields of block size = 1 can currently be projected.
460 
461    The only projection methods currently only support the `DMDA` (2D) and `DMPLEX` (triangles 2D).
462 
463 .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
464 @*/
465 PETSC_EXTERN PetscErrorCode DMSwarmProjectFields(DM dm, PetscInt nfields, const char *fieldnames[], Vec **fields, PetscBool reuse)
466 {
467   DM_Swarm         *swarm = (DM_Swarm *)dm->data;
468   DMSwarmDataField *gfield;
469   DM                celldm;
470   PetscBool         isDA, isPLEX;
471   Vec              *vecs;
472   PetscInt          f, nvecs;
473   PetscInt          project_type = 0;
474 
475   PetscFunctionBegin;
476   DMSWARMPICVALID(dm);
477   PetscCall(DMSwarmGetCellDM(dm, &celldm));
478   PetscCall(PetscMalloc1(nfields, &gfield));
479   nvecs = 0;
480   for (f = 0; f < nfields; f++) {
481     PetscCall(DMSwarmDataBucketGetDMSwarmDataFieldByName(swarm->db, fieldnames[f], &gfield[f]));
482     PetscCheck(gfield[f]->petsc_type == PETSC_REAL, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Projection only valid for fields using a data type = PETSC_REAL");
483     PetscCheck(gfield[f]->bs == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Projection only valid for fields with block size = 1");
484     nvecs += gfield[f]->bs;
485   }
486   if (!reuse) {
487     PetscCall(PetscMalloc1(nvecs, &vecs));
488     for (f = 0; f < nvecs; f++) {
489       PetscCall(DMCreateGlobalVector(celldm, &vecs[f]));
490       PetscCall(PetscObjectSetName((PetscObject)vecs[f], gfield[f]->name));
491     }
492   } else {
493     vecs = *fields;
494   }
495 
496   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
497   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
498   if (isDA) {
499     PetscCall(private_DMSwarmProjectFields_DA(dm, celldm, project_type, nfields, gfield, vecs));
500   } else if (isPLEX) {
501     PetscCall(private_DMSwarmProjectFields_PLEX(dm, celldm, project_type, nfields, gfield, vecs));
502   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
503 
504   PetscCall(PetscFree(gfield));
505   if (!reuse) *fields = vecs;
506   PetscFunctionReturn(PETSC_SUCCESS);
507 }
508 
509 /*@C
510    DMSwarmCreatePointPerCellCount - Count the number of points within all cells in the cell DM
511 
512    Not Collective
513 
514    Input parameter:
515 .  dm - the `DMSWARM`
516 
517    Output parameters:
518 +  ncells - the number of cells in the cell `DM` (optional argument, pass `NULL` to ignore)
519 -  count - array of length ncells containing the number of points per cell
520 
521    Level: beginner
522 
523    Notes:
524    The array count is allocated internally and must be free'd by the user.
525 
526 .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
527 @*/
528 PETSC_EXTERN PetscErrorCode DMSwarmCreatePointPerCellCount(DM dm, PetscInt *ncells, PetscInt **count)
529 {
530   PetscBool isvalid;
531   PetscInt  nel;
532   PetscInt *sum;
533 
534   PetscFunctionBegin;
535   PetscCall(DMSwarmSortGetIsValid(dm, &isvalid));
536   nel = 0;
537   if (isvalid) {
538     PetscInt e;
539 
540     PetscCall(DMSwarmSortGetSizes(dm, &nel, NULL));
541 
542     PetscCall(PetscMalloc1(nel, &sum));
543     for (e = 0; e < nel; e++) PetscCall(DMSwarmSortGetNumberOfPointsPerCell(dm, e, &sum[e]));
544   } else {
545     DM        celldm;
546     PetscBool isda, isplex, isshell;
547     PetscInt  p, npoints;
548     PetscInt *swarm_cellid;
549 
550     /* get the number of cells */
551     PetscCall(DMSwarmGetCellDM(dm, &celldm));
552     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isda));
553     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isplex));
554     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMSHELL, &isshell));
555     if (isda) {
556       PetscInt        _nel, _npe;
557       const PetscInt *_element;
558 
559       PetscCall(DMDAGetElements(celldm, &_nel, &_npe, &_element));
560       nel = _nel;
561       PetscCall(DMDARestoreElements(celldm, &_nel, &_npe, &_element));
562     } else if (isplex) {
563       PetscInt ps, pe;
564 
565       PetscCall(DMPlexGetHeightStratum(celldm, 0, &ps, &pe));
566       nel = pe - ps;
567     } else if (isshell) {
568       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);
569 
570       PetscCall(PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
571       if (method_DMShellGetNumberOfCells) {
572         PetscCall(method_DMShellGetNumberOfCells(celldm, &nel));
573       } else
574         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);");
575     } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");
576 
577     PetscCall(PetscMalloc1(nel, &sum));
578     PetscCall(PetscArrayzero(sum, nel));
579     PetscCall(DMSwarmGetLocalSize(dm, &npoints));
580     PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
581     for (p = 0; p < npoints; p++) {
582       if (swarm_cellid[p] != DMLOCATEPOINT_POINT_NOT_FOUND) sum[swarm_cellid[p]]++;
583     }
584     PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
585   }
586   if (ncells) *ncells = nel;
587   *count = sum;
588   PetscFunctionReturn(PETSC_SUCCESS);
589 }
590 
591 /*@
592   DMSwarmGetNumSpecies - Get the number of particle species
593 
594   Not Collective
595 
596   Input parameter:
597 . dm - the `DMSWARM`
598 
599   Output parameters:
600 . Ns - the number of species
601 
602   Level: intermediate
603 
604 .seealso: `DMSWARM`, `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
605 @*/
606 PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
607 {
608   DM_Swarm *swarm = (DM_Swarm *)sw->data;
609 
610   PetscFunctionBegin;
611   *Ns = swarm->Ns;
612   PetscFunctionReturn(PETSC_SUCCESS);
613 }
614 
615 /*@
616   DMSwarmSetNumSpecies - Set the number of particle species
617 
618   Not Collective
619 
620   Input parameter:
621 + dm - the `DMSWARM`
622 - Ns - the number of species
623 
624   Level: intermediate
625 
626 .seealso: `DMSWARM`, `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
627 @*/
628 PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
629 {
630   DM_Swarm *swarm = (DM_Swarm *)sw->data;
631 
632   PetscFunctionBegin;
633   swarm->Ns = Ns;
634   PetscFunctionReturn(PETSC_SUCCESS);
635 }
636 
637 /*@C
638   DMSwarmGetCoordinateFunction - Get the function setting initial particle positions, if it exists
639 
640   Not Collective
641 
642   Input parameter:
643 . dm - the `DMSWARM`
644 
645   Output Parameter:
646 . coordFunc - the function setting initial particle positions, or `NULL`
647 
648   Level: intermediate
649 
650 .seealso: `DMSWARM`, `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
651 @*/
652 PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFunc *coordFunc)
653 {
654   DM_Swarm *swarm = (DM_Swarm *)sw->data;
655 
656   PetscFunctionBegin;
657   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
658   PetscValidPointer(coordFunc, 2);
659   *coordFunc = swarm->coordFunc;
660   PetscFunctionReturn(PETSC_SUCCESS);
661 }
662 
663 /*@C
664   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions
665 
666   Not Collective
667 
668   Input parameters:
669 + dm - the `DMSWARM`
670 - coordFunc - the function setting initial particle positions
671 
672   Level: intermediate
673 
674 .seealso: `DMSWARM`, `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
675 @*/
676 PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFunc coordFunc)
677 {
678   DM_Swarm *swarm = (DM_Swarm *)sw->data;
679 
680   PetscFunctionBegin;
681   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
682   PetscValidFunction(coordFunc, 2);
683   swarm->coordFunc = coordFunc;
684   PetscFunctionReturn(PETSC_SUCCESS);
685 }
686 
687 /*@C
688   DMSwarmGetCoordinateFunction - Get the function setting initial particle velocities, if it exists
689 
690   Not Collective
691 
692   Input parameter:
693 . dm - the `DMSWARM`
694 
695   Output Parameter:
696 . velFunc - the function setting initial particle velocities, or `NULL`
697 
698   Level: intermediate
699 
700 .seealso: `DMSWARM`, `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
701 @*/
702 PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFunc *velFunc)
703 {
704   DM_Swarm *swarm = (DM_Swarm *)sw->data;
705 
706   PetscFunctionBegin;
707   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
708   PetscValidPointer(velFunc, 2);
709   *velFunc = swarm->velFunc;
710   PetscFunctionReturn(PETSC_SUCCESS);
711 }
712 
713 /*@C
714   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities
715 
716   Not Collective
717 
718   Input parameters:
719 + dm - the `DMSWARM`
720 - coordFunc - the function setting initial particle velocities
721 
722   Level: intermediate
723 
724 .seealso: `DMSWARM`, `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
725 @*/
726 PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFunc velFunc)
727 {
728   DM_Swarm *swarm = (DM_Swarm *)sw->data;
729 
730   PetscFunctionBegin;
731   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
732   PetscValidFunction(velFunc, 2);
733   swarm->velFunc = velFunc;
734   PetscFunctionReturn(PETSC_SUCCESS);
735 }
736 
737 /*@C
738   DMSwarmComputeLocalSize - Compute the local number and distribution of particles based upon a density function
739 
740   Not Collective
741 
742   Input Parameters:
743 + sw      - The `DMSWARM`
744 . N       - The target number of particles
745 - density - The density field for the particle layout, normalized to unity
746 
747   Level: advanced
748 
749   Note:
750   One particle will be created for each species.
751 
752 .seealso: `DMSWARM`, `DMSwarmComputeLocalSizeFromOptions()`
753 @*/
754 PetscErrorCode DMSwarmComputeLocalSize(DM sw, PetscInt N, PetscProbFunc density)
755 {
756   DM               dm;
757   PetscQuadrature  quad;
758   const PetscReal *xq, *wq;
759   PetscReal       *n_int;
760   PetscInt        *npc_s, *cellid, Ni;
761   PetscReal        gmin[3], gmax[3], xi0[3];
762   PetscInt         Ns, cStart, cEnd, c, dim, d, Nq, q, Np = 0, p, s;
763   PetscBool        simplex;
764 
765   PetscFunctionBegin;
766   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
767   PetscCall(DMSwarmGetCellDM(sw, &dm));
768   PetscCall(DMGetDimension(dm, &dim));
769   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
770   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
771   PetscCall(DMPlexIsSimplex(dm, &simplex));
772   PetscCall(DMGetCoordinatesLocalSetUp(dm));
773   if (simplex) PetscCall(PetscDTStroudConicalQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
774   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
775   PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, &xq, &wq));
776   PetscCall(PetscCalloc2(Ns, &n_int, (cEnd - cStart) * Ns, &npc_s));
777   /* Integrate the density function to get the number of particles in each cell */
778   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
779   for (c = 0; c < cEnd - cStart; ++c) {
780     const PetscInt cell = c + cStart;
781     PetscReal      v0[3], J[9], invJ[9], detJ, detJp = 2. / (gmax[0] - gmin[0]), xr[3], den;
782 
783     /*Have to transform quadrature points/weights to cell domain*/
784     PetscCall(DMPlexComputeCellGeometryFEM(dm, cell, NULL, v0, J, invJ, &detJ));
785     PetscCall(PetscArrayzero(n_int, Ns));
786     for (q = 0; q < Nq; ++q) {
787       CoordinatesRefToReal(dim, dim, xi0, v0, J, &xq[q * dim], xr);
788       /* Have to transform mesh to domain of definition of PDF, [-1, 1], and weight PDF by |J|/2 */
789       xr[0] = detJp * (xr[0] - gmin[0]) - 1.;
790 
791       for (s = 0; s < Ns; ++s) {
792         PetscCall(density(xr, NULL, &den));
793         n_int[s] += (detJp * den) * (detJ * wq[q]) / (PetscReal)Ns;
794       }
795     }
796     for (s = 0; s < Ns; ++s) {
797       Ni = N;
798       npc_s[c * Ns + s] += (PetscInt)(Ni * n_int[s]);
799       Np += npc_s[c * Ns + s];
800     }
801   }
802   PetscCall(PetscQuadratureDestroy(&quad));
803   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
804   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
805   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
806     for (s = 0; s < Ns; ++s) {
807       for (q = 0; q < npc_s[c * Ns + s]; ++q, ++p) cellid[p] = c;
808     }
809   }
810   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
811   PetscCall(PetscFree2(n_int, npc_s));
812   PetscFunctionReturn(PETSC_SUCCESS);
813 }
814 
815 /*@
816   DMSwarmComputeLocalSizeFromOptions - Compute the local number and distribution of particles based upon a density function determined by options
817 
818   Not Collective
819 
820   Input Parameter:
821 . sw - The `DMSWARM`
822 
823   Level: advanced
824 
825 .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`
826 @*/
827 PetscErrorCode DMSwarmComputeLocalSizeFromOptions(DM sw)
828 {
829   PetscProbFunc pdf;
830   const char   *prefix;
831   char          funcname[PETSC_MAX_PATH_LEN];
832   PetscInt     *N, Ns, dim, n;
833   PetscBool     flg;
834   PetscMPIInt   size, rank;
835 
836   PetscFunctionBegin;
837   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
838   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
839   PetscCall(PetscCalloc1(size, &N));
840   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
841   n = size;
842   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
843   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
844   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
845   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
846   PetscCall(PetscOptionsString("-dm_swarm_coordinate_function", "Function to determine particle coordinates", "DMSwarmSetCoordinateFunction", funcname, funcname, sizeof(funcname), &flg));
847   PetscOptionsEnd();
848   if (flg) {
849     PetscSimplePointFunc coordFunc;
850 
851     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
852     PetscCall(PetscDLSym(NULL, funcname, (void **)&coordFunc));
853     PetscCheck(coordFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
854     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
855     PetscCall(DMSwarmSetLocalSizes(sw, N[rank] * Ns, 0));
856     PetscCall(DMSwarmSetCoordinateFunction(sw, coordFunc));
857   } else {
858     PetscCall(DMGetDimension(sw, &dim));
859     PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
860     PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_coordinate_density", &pdf, NULL, NULL));
861     PetscCall(DMSwarmComputeLocalSize(sw, N[rank], pdf));
862   }
863   PetscCall(PetscFree(N));
864   PetscFunctionReturn(PETSC_SUCCESS);
865 }
866 
867 /*@
868   DMSwarmInitializeCoordinates - Determine the initial coordinates of particles for a PIC method
869 
870   Not Collective
871 
872   Input Parameter:
873 . sw - The `DMSWARM`
874 
875   Level: advanced
876 
877   Note:
878   Currently, we randomly place particles in their assigned cell
879 
880 .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeVelocities()`
881 @*/
882 PetscErrorCode DMSwarmInitializeCoordinates(DM sw)
883 {
884   PetscSimplePointFunc coordFunc;
885   PetscScalar         *weight;
886   PetscReal           *x;
887   PetscInt            *species;
888   void                *ctx;
889   PetscBool            removePoints = PETSC_TRUE;
890   PetscDataType        dtype;
891   PetscInt             Np, p, Ns, dim, d, bs;
892 
893   PetscFunctionBeginUser;
894   PetscCall(DMGetDimension(sw, &dim));
895   PetscCall(DMSwarmGetLocalSize(sw, &Np));
896   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
897   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));
898 
899   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, &bs, &dtype, (void **)&x));
900   PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight));
901   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
902   if (coordFunc) {
903     PetscCall(DMGetApplicationContext(sw, &ctx));
904     for (p = 0; p < Np; ++p) {
905       PetscScalar X[3];
906 
907       PetscCall((*coordFunc)(dim, 0., NULL, p, X, ctx));
908       for (d = 0; d < dim; ++d) x[p * dim + d] = PetscRealPart(X[d]);
909       weight[p]  = 1.0;
910       species[p] = p % Ns;
911     }
912   } else {
913     DM          dm;
914     PetscRandom rnd;
915     PetscReal   xi0[3];
916     PetscInt    cStart, cEnd, c;
917 
918     PetscCall(DMSwarmGetCellDM(sw, &dm));
919     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
920     PetscCall(DMGetApplicationContext(sw, &ctx));
921 
922     /* Set particle position randomly in cell, set weights to 1 */
923     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd));
924     PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
925     PetscCall(PetscRandomSetFromOptions(rnd));
926     PetscCall(DMSwarmSortGetAccess(sw));
927     for (d = 0; d < dim; ++d) xi0[d] = -1.0;
928     for (c = cStart; c < cEnd; ++c) {
929       PetscReal v0[3], J[9], invJ[9], detJ;
930       PetscInt *pidx, Npc, q;
931 
932       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
933       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ));
934       for (q = 0; q < Npc; ++q) {
935         const PetscInt p = pidx[q];
936         PetscReal      xref[3];
937 
938         for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &xref[d]));
939         CoordinatesRefToReal(dim, dim, xi0, v0, J, xref, &x[p * dim]);
940 
941         weight[p]  = 1.0 / Np;
942         species[p] = p % Ns;
943       }
944       PetscCall(PetscFree(pidx));
945     }
946     PetscCall(PetscRandomDestroy(&rnd));
947     PetscCall(DMSwarmSortRestoreAccess(sw));
948   }
949   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
950   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
951   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
952 
953   PetscCall(DMSwarmMigrate(sw, removePoints));
954   PetscCall(DMLocalizeCoordinates(sw));
955   PetscFunctionReturn(PETSC_SUCCESS);
956 }
957 
958 /*@C
959   DMSwarmInitializeVelocities - Set the initial velocities of particles using a distribution.
960 
961   Collective
962 
963   Input Parameters:
964 + sw      - The `DMSWARM` object
965 . sampler - A function which uniformly samples the velocity PDF
966 - v0      - The velocity scale for nondimensionalization for each species
967 
968   Level: advanced
969 
970   Note:
971   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.
972 
973 .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocitiesFromOptions()`
974 @*/
975 PetscErrorCode DMSwarmInitializeVelocities(DM sw, PetscProbFunc sampler, const PetscReal v0[])
976 {
977   PetscSimplePointFunc velFunc;
978   PetscReal           *v;
979   PetscInt            *species;
980   void                *ctx;
981   PetscInt             dim, Np, p;
982 
983   PetscFunctionBegin;
984   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));
985 
986   PetscCall(DMGetDimension(sw, &dim));
987   PetscCall(DMSwarmGetLocalSize(sw, &Np));
988   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
989   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
990   if (v0[0] == 0.) {
991     PetscCall(PetscArrayzero(v, Np * dim));
992   } else if (velFunc) {
993     PetscCall(DMGetApplicationContext(sw, &ctx));
994     for (p = 0; p < Np; ++p) {
995       PetscInt    s = species[p], d;
996       PetscScalar vel[3];
997 
998       PetscCall((*velFunc)(dim, 0., NULL, p, vel, ctx));
999       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
1000     }
1001   } else {
1002     PetscRandom rnd;
1003 
1004     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
1005     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
1006     PetscCall(PetscRandomSetFromOptions(rnd));
1007 
1008     for (p = 0; p < Np; ++p) {
1009       PetscInt  s = species[p], d;
1010       PetscReal a[3], vel[3];
1011 
1012       for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &a[d]));
1013       PetscCall(sampler(a, NULL, vel));
1014       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * vel[d];
1015     }
1016     PetscCall(PetscRandomDestroy(&rnd));
1017   }
1018   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
1019   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
1020   PetscFunctionReturn(PETSC_SUCCESS);
1021 }
1022 
1023 /*@
1024   DMSwarmInitializeVelocitiesFromOptions - Set the initial velocities of particles using a distribution determined from options.
1025 
1026   Collective
1027 
1028   Input Parameters:
1029 + sw      - The `DMSWARM` object
1030 - v0      - The velocity scale for nondimensionalization for each species
1031 
1032   Level: advanced
1033 
1034 .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
1035 @*/
1036 PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
1037 {
1038   PetscProbFunc sampler;
1039   PetscInt      dim;
1040   const char   *prefix;
1041   char          funcname[PETSC_MAX_PATH_LEN];
1042   PetscBool     flg;
1043 
1044   PetscFunctionBegin;
1045   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
1046   PetscCall(PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg));
1047   PetscOptionsEnd();
1048   if (flg) {
1049     PetscSimplePointFunc velFunc;
1050 
1051     PetscCall(PetscDLSym(NULL, funcname, (void **)&velFunc));
1052     PetscCheck(velFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
1053     PetscCall(DMSwarmSetVelocityFunction(sw, velFunc));
1054   }
1055   PetscCall(DMGetDimension(sw, &dim));
1056   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
1057   PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_velocity_density", NULL, NULL, &sampler));
1058   PetscCall(DMSwarmInitializeVelocities(sw, sampler, v0));
1059   PetscFunctionReturn(PETSC_SUCCESS);
1060 }
1061