xref: /petsc/src/dm/impls/swarm/swarmpic.c (revision a4e35b1925eceef64945ea472b84f2bf06a67b5e)
10e2ec84fSDave May #define PETSCDM_DLL
20e2ec84fSDave May #include <petsc/private/dmswarmimpl.h> /*I   "petscdmswarm.h"   I*/
30e2ec84fSDave May #include <petscsf.h>
4b799feefSDave May #include <petscdmda.h>
5b799feefSDave May #include <petscdmplex.h>
635b38c60SMatthew G. Knepley #include <petscdt.h>
7279f676cSBarry Smith #include "../src/dm/impls/swarm/data_bucket.h"
80e2ec84fSDave May 
935b38c60SMatthew G. Knepley #include <petsc/private/petscfeimpl.h> /* For CoordinatesRefToReal() */
1035b38c60SMatthew G. Knepley 
110e2ec84fSDave May /*
125627991aSBarry Smith  Error checking to ensure the swarm type is correct and that a cell DM has been set
130e2ec84fSDave May */
140e2ec84fSDave May #define DMSWARMPICVALID(dm) \
15f7d195e4SLawrence Mitchell   do { \
160e2ec84fSDave May     DM_Swarm *_swarm = (DM_Swarm *)(dm)->data; \
175f80ce2aSJacob Faibussowitsch     PetscCheck(_swarm->swarm_type == DMSWARM_PIC, PetscObjectComm((PetscObject)(dm)), PETSC_ERR_SUP, "Valid only for DMSwarm-PIC. You must call DMSwarmSetType(dm,DMSWARM_PIC)"); \
1828b400f6SJacob Faibussowitsch     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)"); \
19f7d195e4SLawrence Mitchell   } while (0)
200e2ec84fSDave May 
210e2ec84fSDave May /* Coordinate insertition/addition API */
220e2ec84fSDave May /*@C
2320f4b53cSBarry Smith   DMSwarmSetPointsUniformCoordinates - Set point coordinates in a `DMSWARM` on a regular (ijk) grid
240e2ec84fSDave May 
2520f4b53cSBarry Smith   Collective
260e2ec84fSDave May 
2760225df5SJacob Faibussowitsch   Input Parameters:
2820f4b53cSBarry Smith + dm      - the `DMSWARM`
290e2ec84fSDave May . min     - minimum coordinate values in the x, y, z directions (array of length dim)
300e2ec84fSDave May . max     - maximum coordinate values in the x, y, z directions (array of length dim)
310e2ec84fSDave May . npoints - number of points in each spatial direction (array of length dim)
3220f4b53cSBarry Smith - mode    - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)
330e2ec84fSDave May 
340e2ec84fSDave May   Level: beginner
350e2ec84fSDave May 
360e2ec84fSDave May   Notes:
3720f4b53cSBarry Smith   When using mode = `INSERT_VALUES`, this method will reset the number of particles in the `DMSWARM`
3820f4b53cSBarry Smith   to be npoints[0]*npoints[1] (2D) or npoints[0]*npoints[1]*npoints[2] (3D). When using mode = `ADD_VALUES`,
3920f4b53cSBarry Smith   new points will be appended to any already existing in the `DMSWARM`
400e2ec84fSDave May 
4120f4b53cSBarry Smith .seealso: `DM`, `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
420e2ec84fSDave May @*/
43d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmSetPointsUniformCoordinates(DM dm, PetscReal min[], PetscReal max[], PetscInt npoints[], InsertMode mode)
44d71ae5a4SJacob Faibussowitsch {
450e2ec84fSDave May   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
460e2ec84fSDave May   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
470e2ec84fSDave May   PetscInt           i, j, k, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
480e2ec84fSDave May   Vec                coorlocal;
490e2ec84fSDave May   const PetscScalar *_coor;
500e2ec84fSDave May   DM                 celldm;
510e2ec84fSDave May   PetscReal          dx[3];
522e3d3761SDave May   PetscInt           _npoints[] = {0, 0, 1};
530e2ec84fSDave May   Vec                pos;
540e2ec84fSDave May   PetscScalar       *_pos;
550e2ec84fSDave May   PetscReal         *swarm_coor;
560e2ec84fSDave May   PetscInt          *swarm_cellid;
570e2ec84fSDave May   PetscSF            sfcell = NULL;
580e2ec84fSDave May   const PetscSFNode *LA_sfcell;
590e2ec84fSDave May 
600e2ec84fSDave May   PetscFunctionBegin;
610e2ec84fSDave May   DMSWARMPICVALID(dm);
629566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dm, &celldm));
639566063dSJacob Faibussowitsch   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
649566063dSJacob Faibussowitsch   PetscCall(VecGetSize(coorlocal, &N));
659566063dSJacob Faibussowitsch   PetscCall(VecGetBlockSize(coorlocal, &bs));
660e2ec84fSDave May   N = N / bs;
679566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(coorlocal, &_coor));
680e2ec84fSDave May   for (i = 0; i < N; i++) {
690e2ec84fSDave May     for (b = 0; b < bs; b++) {
70a5f152d1SDave May       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
71a5f152d1SDave May       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
720e2ec84fSDave May     }
730e2ec84fSDave May   }
749566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));
750e2ec84fSDave May 
760e2ec84fSDave May   for (b = 0; b < bs; b++) {
7771844bb8SDave May     if (npoints[b] > 1) {
780e2ec84fSDave May       dx[b] = (max[b] - min[b]) / ((PetscReal)(npoints[b] - 1));
7971844bb8SDave May     } else {
8071844bb8SDave May       dx[b] = 0.0;
8171844bb8SDave May     }
822e3d3761SDave May     _npoints[b] = npoints[b];
830e2ec84fSDave May   }
840e2ec84fSDave May 
850e2ec84fSDave May   /* determine number of points living in the bounding box */
860e2ec84fSDave May   n_estimate = 0;
872e3d3761SDave May   for (k = 0; k < _npoints[2]; k++) {
882e3d3761SDave May     for (j = 0; j < _npoints[1]; j++) {
892e3d3761SDave May       for (i = 0; i < _npoints[0]; i++) {
900e2ec84fSDave May         PetscReal xp[] = {0.0, 0.0, 0.0};
910e2ec84fSDave May         PetscInt  ijk[3];
920e2ec84fSDave May         PetscBool point_inside = PETSC_TRUE;
930e2ec84fSDave May 
940e2ec84fSDave May         ijk[0] = i;
950e2ec84fSDave May         ijk[1] = j;
960e2ec84fSDave May         ijk[2] = k;
97ad540459SPierre Jolivet         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
980e2ec84fSDave May         for (b = 0; b < bs; b++) {
99ad540459SPierre Jolivet           if (xp[b] < gmin[b]) point_inside = PETSC_FALSE;
100ad540459SPierre Jolivet           if (xp[b] > gmax[b]) point_inside = PETSC_FALSE;
1010e2ec84fSDave May         }
102ad540459SPierre Jolivet         if (point_inside) n_estimate++;
1030e2ec84fSDave May       }
1040e2ec84fSDave May     }
1050e2ec84fSDave May   }
1060e2ec84fSDave May 
1070e2ec84fSDave May   /* create candidate list */
1089566063dSJacob Faibussowitsch   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
1099566063dSJacob Faibussowitsch   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
1109566063dSJacob Faibussowitsch   PetscCall(VecSetBlockSize(pos, bs));
1119566063dSJacob Faibussowitsch   PetscCall(VecSetFromOptions(pos));
1129566063dSJacob Faibussowitsch   PetscCall(VecGetArray(pos, &_pos));
1130e2ec84fSDave May 
1140e2ec84fSDave May   n_estimate = 0;
1152e3d3761SDave May   for (k = 0; k < _npoints[2]; k++) {
1162e3d3761SDave May     for (j = 0; j < _npoints[1]; j++) {
1172e3d3761SDave May       for (i = 0; i < _npoints[0]; i++) {
1180e2ec84fSDave May         PetscReal xp[] = {0.0, 0.0, 0.0};
1190e2ec84fSDave May         PetscInt  ijk[3];
1200e2ec84fSDave May         PetscBool point_inside = PETSC_TRUE;
1210e2ec84fSDave May 
1220e2ec84fSDave May         ijk[0] = i;
1230e2ec84fSDave May         ijk[1] = j;
1240e2ec84fSDave May         ijk[2] = k;
125ad540459SPierre Jolivet         for (b = 0; b < bs; b++) xp[b] = min[b] + ijk[b] * dx[b];
1260e2ec84fSDave May         for (b = 0; b < bs; b++) {
127ad540459SPierre Jolivet           if (xp[b] < gmin[b]) point_inside = PETSC_FALSE;
128ad540459SPierre Jolivet           if (xp[b] > gmax[b]) point_inside = PETSC_FALSE;
1290e2ec84fSDave May         }
1300e2ec84fSDave May         if (point_inside) {
131ad540459SPierre Jolivet           for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = xp[b];
1320e2ec84fSDave May           n_estimate++;
1330e2ec84fSDave May         }
1340e2ec84fSDave May       }
1350e2ec84fSDave May     }
1360e2ec84fSDave May   }
1379566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(pos, &_pos));
1380e2ec84fSDave May 
1390e2ec84fSDave May   /* locate points */
1409566063dSJacob Faibussowitsch   PetscCall(DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell));
1419566063dSJacob Faibussowitsch   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
1420e2ec84fSDave May   n_found = 0;
1430e2ec84fSDave May   for (p = 0; p < n_estimate; p++) {
144ad540459SPierre Jolivet     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
1450e2ec84fSDave May   }
1460e2ec84fSDave May 
1470e2ec84fSDave May   /* adjust size */
1480e2ec84fSDave May   if (mode == ADD_VALUES) {
1499566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
1500e2ec84fSDave May     n_new_est = n_curr + n_found;
1519566063dSJacob Faibussowitsch     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
1520e2ec84fSDave May   }
1530e2ec84fSDave May   if (mode == INSERT_VALUES) {
1540e2ec84fSDave May     n_curr    = 0;
1550e2ec84fSDave May     n_new_est = n_found;
1569566063dSJacob Faibussowitsch     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
1570e2ec84fSDave May   }
1580e2ec84fSDave May 
1590e2ec84fSDave May   /* initialize new coords, cell owners, pid */
1609566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(pos, &_coor));
1619566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
1629566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
1630e2ec84fSDave May   n_found = 0;
1640e2ec84fSDave May   for (p = 0; p < n_estimate; p++) {
1650e2ec84fSDave May     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
166ad540459SPierre Jolivet       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
1670e2ec84fSDave May       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
1680e2ec84fSDave May       n_found++;
1690e2ec84fSDave May     }
1700e2ec84fSDave May   }
1719566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
1729566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
1739566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(pos, &_coor));
1740e2ec84fSDave May 
1759566063dSJacob Faibussowitsch   PetscCall(PetscSFDestroy(&sfcell));
1769566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pos));
1773ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
1780e2ec84fSDave May }
1790e2ec84fSDave May 
1800e2ec84fSDave May /*@C
18120f4b53cSBarry Smith   DMSwarmSetPointCoordinates - Set point coordinates in a `DMSWARM` from a user defined list
1820e2ec84fSDave May 
18320f4b53cSBarry Smith   Collective
1840e2ec84fSDave May 
18560225df5SJacob Faibussowitsch   Input Parameters:
18620f4b53cSBarry Smith + dm        - the `DMSWARM`
1870e2ec84fSDave May . npoints   - the number of points to insert
1880e2ec84fSDave May . coor      - the coordinate values
18920f4b53cSBarry Smith . 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
19020f4b53cSBarry Smith - mode      - indicates whether to append points to the swarm (`ADD_VALUES`), or over-ride existing points (`INSERT_VALUES`)
1910e2ec84fSDave May 
1920e2ec84fSDave May   Level: beginner
1930e2ec84fSDave May 
1940e2ec84fSDave May   Notes:
19520f4b53cSBarry Smith   If the user has specified `redundant` as `PETSC_FALSE`, the cell `DM` will attempt to locate the coordinates provided by `coor` within
19620f4b53cSBarry Smith   its sub-domain. If they any values within `coor` are not located in the sub-domain, they will be ignored and will not get
19720f4b53cSBarry Smith   added to the `DMSWARM`.
1980e2ec84fSDave May 
19920f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`, `DMSwarmSetPointsUniformCoordinates()`
2000e2ec84fSDave May @*/
201d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinates(DM dm, PetscInt npoints, PetscReal coor[], PetscBool redundant, InsertMode mode)
202d71ae5a4SJacob Faibussowitsch {
2030e2ec84fSDave May   PetscReal          gmin[] = {PETSC_MAX_REAL, PETSC_MAX_REAL, PETSC_MAX_REAL};
2040e2ec84fSDave May   PetscReal          gmax[] = {PETSC_MIN_REAL, PETSC_MIN_REAL, PETSC_MIN_REAL};
2050e2ec84fSDave May   PetscInt           i, N, bs, b, n_estimate, n_curr, n_new_est, p, n_found;
2060e2ec84fSDave May   Vec                coorlocal;
2070e2ec84fSDave May   const PetscScalar *_coor;
2080e2ec84fSDave May   DM                 celldm;
2090e2ec84fSDave May   Vec                pos;
2100e2ec84fSDave May   PetscScalar       *_pos;
2110e2ec84fSDave May   PetscReal         *swarm_coor;
2120e2ec84fSDave May   PetscInt          *swarm_cellid;
2130e2ec84fSDave May   PetscSF            sfcell = NULL;
2140e2ec84fSDave May   const PetscSFNode *LA_sfcell;
2150e2ec84fSDave May   PetscReal         *my_coor;
2160e2ec84fSDave May   PetscInt           my_npoints;
2170e2ec84fSDave May   PetscMPIInt        rank;
2180e2ec84fSDave May   MPI_Comm           comm;
2190e2ec84fSDave May 
2200e2ec84fSDave May   PetscFunctionBegin;
2210e2ec84fSDave May   DMSWARMPICVALID(dm);
2229566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
2239566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(comm, &rank));
2240e2ec84fSDave May 
2259566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dm, &celldm));
2269566063dSJacob Faibussowitsch   PetscCall(DMGetCoordinatesLocal(celldm, &coorlocal));
2279566063dSJacob Faibussowitsch   PetscCall(VecGetSize(coorlocal, &N));
2289566063dSJacob Faibussowitsch   PetscCall(VecGetBlockSize(coorlocal, &bs));
2290e2ec84fSDave May   N = N / bs;
2309566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(coorlocal, &_coor));
2310e2ec84fSDave May   for (i = 0; i < N; i++) {
2320e2ec84fSDave May     for (b = 0; b < bs; b++) {
233a5f152d1SDave May       gmin[b] = PetscMin(gmin[b], PetscRealPart(_coor[bs * i + b]));
234a5f152d1SDave May       gmax[b] = PetscMax(gmax[b], PetscRealPart(_coor[bs * i + b]));
2350e2ec84fSDave May     }
2360e2ec84fSDave May   }
2379566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(coorlocal, &_coor));
2380e2ec84fSDave May 
2390e2ec84fSDave May   /* broadcast points from rank 0 if requested */
2400e2ec84fSDave May   if (redundant) {
2410e2ec84fSDave May     my_npoints = npoints;
2429566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Bcast(&my_npoints, 1, MPIU_INT, 0, comm));
2430e2ec84fSDave May 
2440e2ec84fSDave May     if (rank > 0) { /* allocate space */
2459566063dSJacob Faibussowitsch       PetscCall(PetscMalloc1(bs * my_npoints, &my_coor));
2460e2ec84fSDave May     } else {
2470e2ec84fSDave May       my_coor = coor;
2480e2ec84fSDave May     }
2499566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Bcast(my_coor, bs * my_npoints, MPIU_REAL, 0, comm));
2500e2ec84fSDave May   } else {
2510e2ec84fSDave May     my_npoints = npoints;
2520e2ec84fSDave May     my_coor    = coor;
2530e2ec84fSDave May   }
2540e2ec84fSDave May 
2550e2ec84fSDave May   /* determine the number of points living in the bounding box */
2560e2ec84fSDave May   n_estimate = 0;
2570e2ec84fSDave May   for (i = 0; i < my_npoints; i++) {
2580e2ec84fSDave May     PetscBool point_inside = PETSC_TRUE;
2590e2ec84fSDave May 
2600e2ec84fSDave May     for (b = 0; b < bs; b++) {
261ad540459SPierre Jolivet       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
262ad540459SPierre Jolivet       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
2630e2ec84fSDave May     }
264ad540459SPierre Jolivet     if (point_inside) n_estimate++;
2650e2ec84fSDave May   }
2660e2ec84fSDave May 
2670e2ec84fSDave May   /* create candidate list */
2689566063dSJacob Faibussowitsch   PetscCall(VecCreate(PETSC_COMM_SELF, &pos));
2699566063dSJacob Faibussowitsch   PetscCall(VecSetSizes(pos, bs * n_estimate, PETSC_DECIDE));
2709566063dSJacob Faibussowitsch   PetscCall(VecSetBlockSize(pos, bs));
2719566063dSJacob Faibussowitsch   PetscCall(VecSetFromOptions(pos));
2729566063dSJacob Faibussowitsch   PetscCall(VecGetArray(pos, &_pos));
2730e2ec84fSDave May 
2740e2ec84fSDave May   n_estimate = 0;
2750e2ec84fSDave May   for (i = 0; i < my_npoints; i++) {
2760e2ec84fSDave May     PetscBool point_inside = PETSC_TRUE;
2770e2ec84fSDave May 
2780e2ec84fSDave May     for (b = 0; b < bs; b++) {
279ad540459SPierre Jolivet       if (my_coor[bs * i + b] < gmin[b]) point_inside = PETSC_FALSE;
280ad540459SPierre Jolivet       if (my_coor[bs * i + b] > gmax[b]) point_inside = PETSC_FALSE;
2810e2ec84fSDave May     }
2820e2ec84fSDave May     if (point_inside) {
283ad540459SPierre Jolivet       for (b = 0; b < bs; b++) _pos[bs * n_estimate + b] = my_coor[bs * i + b];
2840e2ec84fSDave May       n_estimate++;
2850e2ec84fSDave May     }
2860e2ec84fSDave May   }
2879566063dSJacob Faibussowitsch   PetscCall(VecRestoreArray(pos, &_pos));
2880e2ec84fSDave May 
2890e2ec84fSDave May   /* locate points */
2909566063dSJacob Faibussowitsch   PetscCall(DMLocatePoints(celldm, pos, DM_POINTLOCATION_NONE, &sfcell));
2910e2ec84fSDave May 
2929566063dSJacob Faibussowitsch   PetscCall(PetscSFGetGraph(sfcell, NULL, NULL, NULL, &LA_sfcell));
2930e2ec84fSDave May   n_found = 0;
2940e2ec84fSDave May   for (p = 0; p < n_estimate; p++) {
295ad540459SPierre Jolivet     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) n_found++;
2960e2ec84fSDave May   }
2970e2ec84fSDave May 
2980e2ec84fSDave May   /* adjust size */
2990e2ec84fSDave May   if (mode == ADD_VALUES) {
3009566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetLocalSize(dm, &n_curr));
3010e2ec84fSDave May     n_new_est = n_curr + n_found;
3029566063dSJacob Faibussowitsch     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
3030e2ec84fSDave May   }
3040e2ec84fSDave May   if (mode == INSERT_VALUES) {
3050e2ec84fSDave May     n_curr    = 0;
3060e2ec84fSDave May     n_new_est = n_found;
3079566063dSJacob Faibussowitsch     PetscCall(DMSwarmSetLocalSizes(dm, n_new_est, -1));
3080e2ec84fSDave May   }
3090e2ec84fSDave May 
3100e2ec84fSDave May   /* initialize new coords, cell owners, pid */
3119566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(pos, &_coor));
3129566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
3139566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
3140e2ec84fSDave May   n_found = 0;
3150e2ec84fSDave May   for (p = 0; p < n_estimate; p++) {
3160e2ec84fSDave May     if (LA_sfcell[p].index != DMLOCATEPOINT_POINT_NOT_FOUND) {
317ad540459SPierre Jolivet       for (b = 0; b < bs; b++) swarm_coor[bs * (n_curr + n_found) + b] = PetscRealPart(_coor[bs * p + b]);
3180e2ec84fSDave May       swarm_cellid[n_curr + n_found] = LA_sfcell[p].index;
3190e2ec84fSDave May       n_found++;
3200e2ec84fSDave May     }
3210e2ec84fSDave May   }
3229566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
3239566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_coor, NULL, NULL, (void **)&swarm_coor));
3249566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(pos, &_coor));
3250e2ec84fSDave May 
3260e2ec84fSDave May   if (redundant) {
32748a46eb9SPierre Jolivet     if (rank > 0) PetscCall(PetscFree(my_coor));
3280e2ec84fSDave May   }
3299566063dSJacob Faibussowitsch   PetscCall(PetscSFDestroy(&sfcell));
3309566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&pos));
3313ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3320e2ec84fSDave May }
3330e2ec84fSDave May 
3340e2ec84fSDave May extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_DA(DM, DM, DMSwarmPICLayoutType, PetscInt);
3350e2ec84fSDave May extern PetscErrorCode private_DMSwarmInsertPointsUsingCellDM_PLEX(DM, DM, DMSwarmPICLayoutType, PetscInt);
3360e2ec84fSDave May 
3370e2ec84fSDave May /*@C
3380e2ec84fSDave May   DMSwarmInsertPointsUsingCellDM - Insert point coordinates within each cell
3390e2ec84fSDave May 
34020f4b53cSBarry Smith   Not Collective
3410e2ec84fSDave May 
34260225df5SJacob Faibussowitsch   Input Parameters:
34320f4b53cSBarry Smith + dm          - the `DMSWARM`
34420f4b53cSBarry Smith . layout_type - method used to fill each cell with the cell `DM`
3450e2ec84fSDave May - fill_param  - parameter controlling how many points per cell are added (the meaning of this parameter is dependent on the layout type)
3460e2ec84fSDave May 
3470e2ec84fSDave May   Level: beginner
3480e2ec84fSDave May 
3490e2ec84fSDave May   Notes:
35020f4b53cSBarry Smith   The insert method will reset any previous defined points within the `DMSWARM`.
351e7af74c8SDave May 
35220f4b53cSBarry Smith   When using a `DMDA` both 2D and 3D are supported for all layout types provided you are using `DMDA_ELEMENT_Q1`.
353e7af74c8SDave May 
354*a4e35b19SJacob Faibussowitsch   When using a `DMPLEX` the following case are supported\:
35520f4b53cSBarry Smith .vb
356ea3d7275SDave May    (i) DMSWARMPIC_LAYOUT_REGULAR: 2D (triangle),
357ea3d7275SDave May    (ii) DMSWARMPIC_LAYOUT_GAUSS: 2D and 3D provided the cell is a tri/tet or a quad/hex,
358ea3d7275SDave May    (iii) DMSWARMPIC_LAYOUT_SUBDIVISION: 2D and 3D for quad/hex and 2D tri.
35920f4b53cSBarry Smith .ve
3600e2ec84fSDave May 
36120f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmPICLayoutType`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
3620e2ec84fSDave May @*/
363d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmInsertPointsUsingCellDM(DM dm, DMSwarmPICLayoutType layout_type, PetscInt fill_param)
364d71ae5a4SJacob Faibussowitsch {
3650e2ec84fSDave May   DM        celldm;
3660e2ec84fSDave May   PetscBool isDA, isPLEX;
3670e2ec84fSDave May 
3680e2ec84fSDave May   PetscFunctionBegin;
3690e2ec84fSDave May   DMSWARMPICVALID(dm);
3709566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dm, &celldm));
3719566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
3729566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
3730e2ec84fSDave May   if (isDA) {
3749566063dSJacob Faibussowitsch     PetscCall(private_DMSwarmInsertPointsUsingCellDM_DA(dm, celldm, layout_type, fill_param));
3750e2ec84fSDave May   } else if (isPLEX) {
3769566063dSJacob Faibussowitsch     PetscCall(private_DMSwarmInsertPointsUsingCellDM_PLEX(dm, celldm, layout_type, fill_param));
3770e2ec84fSDave May   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
3783ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3790e2ec84fSDave May }
3800e2ec84fSDave May 
381431382f2SDave May extern PetscErrorCode private_DMSwarmSetPointCoordinatesCellwise_PLEX(DM, DM, PetscInt, PetscReal *);
382431382f2SDave May 
383431382f2SDave May /*@C
384431382f2SDave May   DMSwarmSetPointCoordinatesCellwise - Insert point coordinates (defined over the reference cell) within each cell
385431382f2SDave May 
38620f4b53cSBarry Smith   Not Collective
387431382f2SDave May 
38860225df5SJacob Faibussowitsch   Input Parameters:
38920f4b53cSBarry Smith + dm      - the `DMSWARM`
390431382f2SDave May . npoints - the number of points to insert in each cell
391431382f2SDave May - xi      - the coordinates (defined in the local coordinate system for each cell) to insert
392431382f2SDave May 
393431382f2SDave May   Level: beginner
394431382f2SDave May 
395431382f2SDave May   Notes:
39620f4b53cSBarry Smith   The method will reset any previous defined points within the `DMSWARM`.
39720f4b53cSBarry Smith   Only supported for `DMPLEX`. If you are using a `DMDA` it is recommended to either use
39820f4b53cSBarry Smith   `DMSwarmInsertPointsUsingCellDM()`, or extract and set the coordinates yourself the following code
39920f4b53cSBarry Smith .vb
40020f4b53cSBarry Smith     PetscReal *coor;
40120f4b53cSBarry Smith     DMSwarmGetField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
40220f4b53cSBarry Smith     // user code to define the coordinates here
40320f4b53cSBarry Smith     DMSwarmRestoreField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
40420f4b53cSBarry Smith .ve
405e7af74c8SDave May 
40620f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
407431382f2SDave May @*/
408d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
409d71ae5a4SJacob Faibussowitsch {
410431382f2SDave May   DM        celldm;
411431382f2SDave May   PetscBool isDA, isPLEX;
412431382f2SDave May 
4130e2ec84fSDave May   PetscFunctionBegin;
414431382f2SDave May   DMSWARMPICVALID(dm);
4159566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dm, &celldm));
4169566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
4179566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
41828b400f6SJacob Faibussowitsch   PetscCheck(!isDA, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMPLEX. Recommended you use DMSwarmInsertPointsUsingCellDM()");
419f7d195e4SLawrence Mitchell   if (isPLEX) {
4209566063dSJacob Faibussowitsch     PetscCall(private_DMSwarmSetPointCoordinatesCellwise_PLEX(dm, celldm, npoints, xi));
421431382f2SDave May   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
4223ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4230e2ec84fSDave May }
424431382f2SDave May 
4250e2ec84fSDave May /* Field projection API */
42677048351SPatrick Sanan extern PetscErrorCode private_DMSwarmProjectFields_DA(DM swarm, DM celldm, PetscInt project_type, PetscInt nfields, DMSwarmDataField dfield[], Vec vecs[]);
42777048351SPatrick Sanan extern PetscErrorCode private_DMSwarmProjectFields_PLEX(DM swarm, DM celldm, PetscInt project_type, PetscInt nfields, DMSwarmDataField dfield[], Vec vecs[]);
42894f7d2dcSDave May 
42994f7d2dcSDave May /*@C
43020f4b53cSBarry Smith   DMSwarmProjectFields - Project a set of swarm fields onto the cell `DM`
43194f7d2dcSDave May 
43220f4b53cSBarry Smith   Collective
43394f7d2dcSDave May 
43460225df5SJacob Faibussowitsch   Input Parameters:
43520f4b53cSBarry Smith + dm         - the `DMSWARM`
43694f7d2dcSDave May . nfields    - the number of swarm fields to project
43794f7d2dcSDave May . fieldnames - the textual names of the swarm fields to project
43894f7d2dcSDave May . fields     - an array of Vec's of length nfields
43994f7d2dcSDave May - reuse      - flag indicating whether the array and contents of fields should be re-used or internally allocated
44094f7d2dcSDave May 
44120f4b53cSBarry Smith   Level: beginner
44220f4b53cSBarry Smith 
44320f4b53cSBarry Smith   Notes:
44494f7d2dcSDave May   Currently, the only available projection method consists of
44520f4b53cSBarry Smith .vb
44694f7d2dcSDave May      phi_i = \sum_{p=0}^{np} N_i(x_p) phi_p dJ / \sum_{p=0}^{np} N_i(x_p) dJ
44794f7d2dcSDave May    where phi_p is the swarm field at point p,
44894f7d2dcSDave May      N_i() is the cell DM basis function at vertex i,
44994f7d2dcSDave May      dJ is the determinant of the cell Jacobian and
45094f7d2dcSDave May      phi_i is the projected vertex value of the field phi.
45120f4b53cSBarry Smith .ve
45294f7d2dcSDave May 
45320f4b53cSBarry Smith   If `reuse` is `PETSC_FALSE`, this function will allocate the array of `Vec`'s, and each individual `Vec`.
45420f4b53cSBarry Smith   The user is responsible for destroying both the array and the individual `Vec` objects.
45594f7d2dcSDave May 
45620f4b53cSBarry Smith   Only swarm fields registered with data type of `PETSC_REAL` can be projected onto the cell `DM`.
457e7af74c8SDave May 
458e7af74c8SDave May   Only swarm fields of block size = 1 can currently be projected.
459e7af74c8SDave May 
46020f4b53cSBarry Smith   The only projection methods currently only support the `DMDA` (2D) and `DMPLEX` (triangles 2D).
46194f7d2dcSDave May 
46220f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
46394f7d2dcSDave May @*/
464d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmProjectFields(DM dm, PetscInt nfields, const char *fieldnames[], Vec **fields, PetscBool reuse)
465d71ae5a4SJacob Faibussowitsch {
46694f7d2dcSDave May   DM_Swarm         *swarm = (DM_Swarm *)dm->data;
46777048351SPatrick Sanan   DMSwarmDataField *gfield;
46894f7d2dcSDave May   DM                celldm;
46994f7d2dcSDave May   PetscBool         isDA, isPLEX;
47094f7d2dcSDave May   Vec              *vecs;
47194f7d2dcSDave May   PetscInt          f, nvecs;
47294f7d2dcSDave May   PetscInt          project_type = 0;
47394f7d2dcSDave May 
4740e2ec84fSDave May   PetscFunctionBegin;
47594f7d2dcSDave May   DMSWARMPICVALID(dm);
4769566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dm, &celldm));
4779566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(nfields, &gfield));
47894f7d2dcSDave May   nvecs = 0;
47994f7d2dcSDave May   for (f = 0; f < nfields; f++) {
4809566063dSJacob Faibussowitsch     PetscCall(DMSwarmDataBucketGetDMSwarmDataFieldByName(swarm->db, fieldnames[f], &gfield[f]));
4815f80ce2aSJacob Faibussowitsch     PetscCheck(gfield[f]->petsc_type == PETSC_REAL, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Projection only valid for fields using a data type = PETSC_REAL");
4825f80ce2aSJacob Faibussowitsch     PetscCheck(gfield[f]->bs == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Projection only valid for fields with block size = 1");
48394f7d2dcSDave May     nvecs += gfield[f]->bs;
48494f7d2dcSDave May   }
48594f7d2dcSDave May   if (!reuse) {
4869566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(nvecs, &vecs));
48794f7d2dcSDave May     for (f = 0; f < nvecs; f++) {
4889566063dSJacob Faibussowitsch       PetscCall(DMCreateGlobalVector(celldm, &vecs[f]));
4899566063dSJacob Faibussowitsch       PetscCall(PetscObjectSetName((PetscObject)vecs[f], gfield[f]->name));
49094f7d2dcSDave May     }
49194f7d2dcSDave May   } else {
49294f7d2dcSDave May     vecs = *fields;
49394f7d2dcSDave May   }
49494f7d2dcSDave May 
4959566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
4969566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
49794f7d2dcSDave May   if (isDA) {
4989566063dSJacob Faibussowitsch     PetscCall(private_DMSwarmProjectFields_DA(dm, celldm, project_type, nfields, gfield, vecs));
49994f7d2dcSDave May   } else if (isPLEX) {
5009566063dSJacob Faibussowitsch     PetscCall(private_DMSwarmProjectFields_PLEX(dm, celldm, project_type, nfields, gfield, vecs));
50194f7d2dcSDave May   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
50294f7d2dcSDave May 
5039566063dSJacob Faibussowitsch   PetscCall(PetscFree(gfield));
5045f80ce2aSJacob Faibussowitsch   if (!reuse) *fields = vecs;
5053ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5060e2ec84fSDave May }
5070e2ec84fSDave May 
5080e2ec84fSDave May /*@C
509b799feefSDave May   DMSwarmCreatePointPerCellCount - Count the number of points within all cells in the cell DM
5100e2ec84fSDave May 
51120f4b53cSBarry Smith   Not Collective
5120e2ec84fSDave May 
51360225df5SJacob Faibussowitsch   Input Parameter:
51420f4b53cSBarry Smith . dm - the `DMSWARM`
5150e2ec84fSDave May 
51660225df5SJacob Faibussowitsch   Output Parameters:
51720f4b53cSBarry Smith + ncells - the number of cells in the cell `DM` (optional argument, pass `NULL` to ignore)
518b799feefSDave May - count  - array of length ncells containing the number of points per cell
5190e2ec84fSDave May 
5200e2ec84fSDave May   Level: beginner
5210e2ec84fSDave May 
5220e2ec84fSDave May   Notes:
5230e2ec84fSDave May   The array count is allocated internally and must be free'd by the user.
5240e2ec84fSDave May 
52520f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
5260e2ec84fSDave May @*/
527d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmCreatePointPerCellCount(DM dm, PetscInt *ncells, PetscInt **count)
528d71ae5a4SJacob Faibussowitsch {
529b799feefSDave May   PetscBool isvalid;
530b799feefSDave May   PetscInt  nel;
531b799feefSDave May   PetscInt *sum;
532b799feefSDave May 
5330e2ec84fSDave May   PetscFunctionBegin;
5349566063dSJacob Faibussowitsch   PetscCall(DMSwarmSortGetIsValid(dm, &isvalid));
535b799feefSDave May   nel = 0;
536b799feefSDave May   if (isvalid) {
537b799feefSDave May     PetscInt e;
538b799feefSDave May 
5399566063dSJacob Faibussowitsch     PetscCall(DMSwarmSortGetSizes(dm, &nel, NULL));
540b799feefSDave May 
5419566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(nel, &sum));
54248a46eb9SPierre Jolivet     for (e = 0; e < nel; e++) PetscCall(DMSwarmSortGetNumberOfPointsPerCell(dm, e, &sum[e]));
543b799feefSDave May   } else {
544b799feefSDave May     DM        celldm;
545b799feefSDave May     PetscBool isda, isplex, isshell;
546b799feefSDave May     PetscInt  p, npoints;
547b799feefSDave May     PetscInt *swarm_cellid;
548b799feefSDave May 
549b799feefSDave May     /* get the number of cells */
5509566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetCellDM(dm, &celldm));
5519566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isda));
5529566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isplex));
5539566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMSHELL, &isshell));
554b799feefSDave May     if (isda) {
555b799feefSDave May       PetscInt        _nel, _npe;
556b799feefSDave May       const PetscInt *_element;
557b799feefSDave May 
5589566063dSJacob Faibussowitsch       PetscCall(DMDAGetElements(celldm, &_nel, &_npe, &_element));
559b799feefSDave May       nel = _nel;
5609566063dSJacob Faibussowitsch       PetscCall(DMDARestoreElements(celldm, &_nel, &_npe, &_element));
561b799feefSDave May     } else if (isplex) {
562b799feefSDave May       PetscInt ps, pe;
563b799feefSDave May 
5649566063dSJacob Faibussowitsch       PetscCall(DMPlexGetHeightStratum(celldm, 0, &ps, &pe));
565b799feefSDave May       nel = pe - ps;
566b799feefSDave May     } else if (isshell) {
567b799feefSDave May       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);
568b799feefSDave May 
5699566063dSJacob Faibussowitsch       PetscCall(PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
570b799feefSDave May       if (method_DMShellGetNumberOfCells) {
5719566063dSJacob Faibussowitsch         PetscCall(method_DMShellGetNumberOfCells(celldm, &nel));
5729371c9d4SSatish Balay       } else
5739371c9d4SSatish Balay         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);");
574b799feefSDave May     } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot determine the number of cells for a DM not of type DA, PLEX or SHELL");
575b799feefSDave May 
5769566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(nel, &sum));
5779566063dSJacob Faibussowitsch     PetscCall(PetscArrayzero(sum, nel));
5789566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetLocalSize(dm, &npoints));
5799566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
580b799feefSDave May     for (p = 0; p < npoints; p++) {
581ad540459SPierre Jolivet       if (swarm_cellid[p] != DMLOCATEPOINT_POINT_NOT_FOUND) sum[swarm_cellid[p]]++;
582b799feefSDave May     }
5839566063dSJacob Faibussowitsch     PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
584b799feefSDave May   }
585ad540459SPierre Jolivet   if (ncells) *ncells = nel;
586b799feefSDave May   *count = sum;
5873ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5880e2ec84fSDave May }
58935b38c60SMatthew G. Knepley 
59035b38c60SMatthew G. Knepley /*@
59135b38c60SMatthew G. Knepley   DMSwarmGetNumSpecies - Get the number of particle species
59235b38c60SMatthew G. Knepley 
59320f4b53cSBarry Smith   Not Collective
59435b38c60SMatthew G. Knepley 
59560225df5SJacob Faibussowitsch   Input Parameter:
59660225df5SJacob Faibussowitsch . sw - the `DMSWARM`
59735b38c60SMatthew G. Knepley 
59860225df5SJacob Faibussowitsch   Output Parameters:
59935b38c60SMatthew G. Knepley . Ns - the number of species
60035b38c60SMatthew G. Knepley 
60135b38c60SMatthew G. Knepley   Level: intermediate
60235b38c60SMatthew G. Knepley 
60320f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
60435b38c60SMatthew G. Knepley @*/
605d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
606d71ae5a4SJacob Faibussowitsch {
60735b38c60SMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
60835b38c60SMatthew G. Knepley 
60935b38c60SMatthew G. Knepley   PetscFunctionBegin;
61035b38c60SMatthew G. Knepley   *Ns = swarm->Ns;
6113ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
61235b38c60SMatthew G. Knepley }
61335b38c60SMatthew G. Knepley 
61435b38c60SMatthew G. Knepley /*@
61535b38c60SMatthew G. Knepley   DMSwarmSetNumSpecies - Set the number of particle species
61635b38c60SMatthew G. Knepley 
61720f4b53cSBarry Smith   Not Collective
61835b38c60SMatthew G. Knepley 
61960225df5SJacob Faibussowitsch   Input Parameters:
62060225df5SJacob Faibussowitsch + sw - the `DMSWARM`
62135b38c60SMatthew G. Knepley - Ns - the number of species
62235b38c60SMatthew G. Knepley 
62335b38c60SMatthew G. Knepley   Level: intermediate
62435b38c60SMatthew G. Knepley 
62520f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
62635b38c60SMatthew G. Knepley @*/
627d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
628d71ae5a4SJacob Faibussowitsch {
62935b38c60SMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
63035b38c60SMatthew G. Knepley 
63135b38c60SMatthew G. Knepley   PetscFunctionBegin;
63235b38c60SMatthew G. Knepley   swarm->Ns = Ns;
6333ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
63435b38c60SMatthew G. Knepley }
63535b38c60SMatthew G. Knepley 
63635b38c60SMatthew G. Knepley /*@C
6376c5a79ebSMatthew G. Knepley   DMSwarmGetCoordinateFunction - Get the function setting initial particle positions, if it exists
6386c5a79ebSMatthew G. Knepley 
63920f4b53cSBarry Smith   Not Collective
6406c5a79ebSMatthew G. Knepley 
64160225df5SJacob Faibussowitsch   Input Parameter:
64260225df5SJacob Faibussowitsch . sw - the `DMSWARM`
6436c5a79ebSMatthew G. Knepley 
6446c5a79ebSMatthew G. Knepley   Output Parameter:
64520f4b53cSBarry Smith . coordFunc - the function setting initial particle positions, or `NULL`
6466c5a79ebSMatthew G. Knepley 
6476c5a79ebSMatthew G. Knepley   Level: intermediate
6486c5a79ebSMatthew G. Knepley 
64920f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
6506c5a79ebSMatthew G. Knepley @*/
651d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFunc *coordFunc)
652d71ae5a4SJacob Faibussowitsch {
6536c5a79ebSMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
6546c5a79ebSMatthew G. Knepley 
6556c5a79ebSMatthew G. Knepley   PetscFunctionBegin;
6566c5a79ebSMatthew G. Knepley   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
6574f572ea9SToby Isaac   PetscAssertPointer(coordFunc, 2);
6586c5a79ebSMatthew G. Knepley   *coordFunc = swarm->coordFunc;
6593ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
6606c5a79ebSMatthew G. Knepley }
6616c5a79ebSMatthew G. Knepley 
6626c5a79ebSMatthew G. Knepley /*@C
6636c5a79ebSMatthew G. Knepley   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions
6646c5a79ebSMatthew G. Knepley 
66520f4b53cSBarry Smith   Not Collective
6666c5a79ebSMatthew G. Knepley 
66760225df5SJacob Faibussowitsch   Input Parameters:
66860225df5SJacob Faibussowitsch + sw        - the `DMSWARM`
6696c5a79ebSMatthew G. Knepley - coordFunc - the function setting initial particle positions
6706c5a79ebSMatthew G. Knepley 
6716c5a79ebSMatthew G. Knepley   Level: intermediate
6726c5a79ebSMatthew G. Knepley 
67320f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
6746c5a79ebSMatthew G. Knepley @*/
675d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFunc coordFunc)
676d71ae5a4SJacob Faibussowitsch {
6776c5a79ebSMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
6786c5a79ebSMatthew G. Knepley 
6796c5a79ebSMatthew G. Knepley   PetscFunctionBegin;
6806c5a79ebSMatthew G. Knepley   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
6816c5a79ebSMatthew G. Knepley   PetscValidFunction(coordFunc, 2);
6826c5a79ebSMatthew G. Knepley   swarm->coordFunc = coordFunc;
6833ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
6846c5a79ebSMatthew G. Knepley }
6856c5a79ebSMatthew G. Knepley 
6866c5a79ebSMatthew G. Knepley /*@C
68760225df5SJacob Faibussowitsch   DMSwarmGetVelocityFunction - Get the function setting initial particle velocities, if it exists
6886c5a79ebSMatthew G. Knepley 
68920f4b53cSBarry Smith   Not Collective
6906c5a79ebSMatthew G. Knepley 
69160225df5SJacob Faibussowitsch   Input Parameter:
69260225df5SJacob Faibussowitsch . sw - the `DMSWARM`
6936c5a79ebSMatthew G. Knepley 
6946c5a79ebSMatthew G. Knepley   Output Parameter:
69520f4b53cSBarry Smith . velFunc - the function setting initial particle velocities, or `NULL`
6966c5a79ebSMatthew G. Knepley 
6976c5a79ebSMatthew G. Knepley   Level: intermediate
6986c5a79ebSMatthew G. Knepley 
69920f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
7006c5a79ebSMatthew G. Knepley @*/
701d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFunc *velFunc)
702d71ae5a4SJacob Faibussowitsch {
7036c5a79ebSMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
7046c5a79ebSMatthew G. Knepley 
7056c5a79ebSMatthew G. Knepley   PetscFunctionBegin;
7066c5a79ebSMatthew G. Knepley   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
7074f572ea9SToby Isaac   PetscAssertPointer(velFunc, 2);
7086c5a79ebSMatthew G. Knepley   *velFunc = swarm->velFunc;
7093ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7106c5a79ebSMatthew G. Knepley }
7116c5a79ebSMatthew G. Knepley 
7126c5a79ebSMatthew G. Knepley /*@C
7136c5a79ebSMatthew G. Knepley   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities
7146c5a79ebSMatthew G. Knepley 
71520f4b53cSBarry Smith   Not Collective
7166c5a79ebSMatthew G. Knepley 
71760225df5SJacob Faibussowitsch   Input Parameters:
718*a4e35b19SJacob Faibussowitsch + sw      - the `DMSWARM`
719*a4e35b19SJacob Faibussowitsch - velFunc - the function setting initial particle velocities
7206c5a79ebSMatthew G. Knepley 
7216c5a79ebSMatthew G. Knepley   Level: intermediate
7226c5a79ebSMatthew G. Knepley 
72320f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
7246c5a79ebSMatthew G. Knepley @*/
725d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFunc velFunc)
726d71ae5a4SJacob Faibussowitsch {
7276c5a79ebSMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
7286c5a79ebSMatthew G. Knepley 
7296c5a79ebSMatthew G. Knepley   PetscFunctionBegin;
7306c5a79ebSMatthew G. Knepley   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
7316c5a79ebSMatthew G. Knepley   PetscValidFunction(velFunc, 2);
7326c5a79ebSMatthew G. Knepley   swarm->velFunc = velFunc;
7333ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7346c5a79ebSMatthew G. Knepley }
7356c5a79ebSMatthew G. Knepley 
7366c5a79ebSMatthew G. Knepley /*@C
73735b38c60SMatthew G. Knepley   DMSwarmComputeLocalSize - Compute the local number and distribution of particles based upon a density function
73835b38c60SMatthew G. Knepley 
73920f4b53cSBarry Smith   Not Collective
74035b38c60SMatthew G. Knepley 
74135b38c60SMatthew G. Knepley   Input Parameters:
74220f4b53cSBarry Smith + sw      - The `DMSWARM`
74335b38c60SMatthew G. Knepley . N       - The target number of particles
74435b38c60SMatthew G. Knepley - density - The density field for the particle layout, normalized to unity
74535b38c60SMatthew G. Knepley 
74635b38c60SMatthew G. Knepley   Level: advanced
74735b38c60SMatthew G. Knepley 
74820f4b53cSBarry Smith   Note:
74920f4b53cSBarry Smith   One particle will be created for each species.
75020f4b53cSBarry Smith 
75120f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSizeFromOptions()`
75235b38c60SMatthew G. Knepley @*/
753d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmComputeLocalSize(DM sw, PetscInt N, PetscProbFunc density)
754d71ae5a4SJacob Faibussowitsch {
75535b38c60SMatthew G. Knepley   DM               dm;
75635b38c60SMatthew G. Knepley   PetscQuadrature  quad;
75735b38c60SMatthew G. Knepley   const PetscReal *xq, *wq;
758ea7032a0SMatthew G. Knepley   PetscReal       *n_int;
759ea7032a0SMatthew G. Knepley   PetscInt        *npc_s, *cellid, Ni;
760ea7032a0SMatthew G. Knepley   PetscReal        gmin[3], gmax[3], xi0[3];
761ea7032a0SMatthew G. Knepley   PetscInt         Ns, cStart, cEnd, c, dim, d, Nq, q, Np = 0, p, s;
76235b38c60SMatthew G. Knepley   PetscBool        simplex;
76335b38c60SMatthew G. Knepley 
76435b38c60SMatthew G. Knepley   PetscFunctionBegin;
7659566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
7669566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(sw, &dm));
7679566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(dm, &dim));
768ea7032a0SMatthew G. Knepley   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
7699566063dSJacob Faibussowitsch   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
7709566063dSJacob Faibussowitsch   PetscCall(DMPlexIsSimplex(dm, &simplex));
7716858538eSMatthew G. Knepley   PetscCall(DMGetCoordinatesLocalSetUp(dm));
7729566063dSJacob Faibussowitsch   if (simplex) PetscCall(PetscDTStroudConicalQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
7739566063dSJacob Faibussowitsch   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
7749566063dSJacob Faibussowitsch   PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, &xq, &wq));
775ea7032a0SMatthew G. Knepley   PetscCall(PetscCalloc2(Ns, &n_int, (cEnd - cStart) * Ns, &npc_s));
77635b38c60SMatthew G. Knepley   /* Integrate the density function to get the number of particles in each cell */
77735b38c60SMatthew G. Knepley   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
77835b38c60SMatthew G. Knepley   for (c = 0; c < cEnd - cStart; ++c) {
77935b38c60SMatthew G. Knepley     const PetscInt cell = c + cStart;
780ea7032a0SMatthew G. Knepley     PetscReal      v0[3], J[9], invJ[9], detJ, detJp = 2. / (gmax[0] - gmin[0]), xr[3], den;
78135b38c60SMatthew G. Knepley 
782ea7032a0SMatthew G. Knepley     /*Have to transform quadrature points/weights to cell domain*/
7839566063dSJacob Faibussowitsch     PetscCall(DMPlexComputeCellGeometryFEM(dm, cell, NULL, v0, J, invJ, &detJ));
784ea7032a0SMatthew G. Knepley     PetscCall(PetscArrayzero(n_int, Ns));
78535b38c60SMatthew G. Knepley     for (q = 0; q < Nq; ++q) {
78635b38c60SMatthew G. Knepley       CoordinatesRefToReal(dim, dim, xi0, v0, J, &xq[q * dim], xr);
787ea7032a0SMatthew G. Knepley       /* Have to transform mesh to domain of definition of PDF, [-1, 1], and weight PDF by |J|/2 */
788ea7032a0SMatthew G. Knepley       xr[0] = detJp * (xr[0] - gmin[0]) - 1.;
789ea7032a0SMatthew G. Knepley 
790ea7032a0SMatthew G. Knepley       for (s = 0; s < Ns; ++s) {
791ea7032a0SMatthew G. Knepley         PetscCall(density(xr, NULL, &den));
792ea7032a0SMatthew G. Knepley         n_int[s] += (detJp * den) * (detJ * wq[q]) / (PetscReal)Ns;
79335b38c60SMatthew G. Knepley       }
794ea7032a0SMatthew G. Knepley     }
795ea7032a0SMatthew G. Knepley     for (s = 0; s < Ns; ++s) {
796ea7032a0SMatthew G. Knepley       Ni = N;
797ea7032a0SMatthew G. Knepley       npc_s[c * Ns + s] += (PetscInt)(Ni * n_int[s]);
798ea7032a0SMatthew G. Knepley       Np += npc_s[c * Ns + s];
799ea7032a0SMatthew G. Knepley     }
80035b38c60SMatthew G. Knepley   }
8019566063dSJacob Faibussowitsch   PetscCall(PetscQuadratureDestroy(&quad));
8029566063dSJacob Faibussowitsch   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
8039566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
80435b38c60SMatthew G. Knepley   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
805ea7032a0SMatthew G. Knepley     for (s = 0; s < Ns; ++s) {
806ea7032a0SMatthew G. Knepley       for (q = 0; q < npc_s[c * Ns + s]; ++q, ++p) cellid[p] = c;
807ea7032a0SMatthew G. Knepley     }
80835b38c60SMatthew G. Knepley   }
8099566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
810ea7032a0SMatthew G. Knepley   PetscCall(PetscFree2(n_int, npc_s));
8113ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
81235b38c60SMatthew G. Knepley }
81335b38c60SMatthew G. Knepley 
81435b38c60SMatthew G. Knepley /*@
81535b38c60SMatthew G. Knepley   DMSwarmComputeLocalSizeFromOptions - Compute the local number and distribution of particles based upon a density function determined by options
81635b38c60SMatthew G. Knepley 
81720f4b53cSBarry Smith   Not Collective
81835b38c60SMatthew G. Knepley 
8192fe279fdSBarry Smith   Input Parameter:
8202fe279fdSBarry Smith . sw - The `DMSWARM`
82135b38c60SMatthew G. Knepley 
82235b38c60SMatthew G. Knepley   Level: advanced
82335b38c60SMatthew G. Knepley 
82420f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`
82535b38c60SMatthew G. Knepley @*/
826d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmComputeLocalSizeFromOptions(DM sw)
827d71ae5a4SJacob Faibussowitsch {
82835b38c60SMatthew G. Knepley   PetscProbFunc pdf;
82935b38c60SMatthew G. Knepley   const char   *prefix;
8306c5a79ebSMatthew G. Knepley   char          funcname[PETSC_MAX_PATH_LEN];
8316c5a79ebSMatthew G. Knepley   PetscInt     *N, Ns, dim, n;
8326c5a79ebSMatthew G. Knepley   PetscBool     flg;
8336c5a79ebSMatthew G. Knepley   PetscMPIInt   size, rank;
83435b38c60SMatthew G. Knepley 
83535b38c60SMatthew G. Knepley   PetscFunctionBegin;
8366c5a79ebSMatthew G. Knepley   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
8376c5a79ebSMatthew G. Knepley   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
8386c5a79ebSMatthew G. Knepley   PetscCall(PetscCalloc1(size, &N));
839d0609cedSBarry Smith   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
8406c5a79ebSMatthew G. Knepley   n = size;
8416c5a79ebSMatthew G. Knepley   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
8426c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
8439566063dSJacob Faibussowitsch   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
8449566063dSJacob Faibussowitsch   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
8456c5a79ebSMatthew G. Knepley   PetscCall(PetscOptionsString("-dm_swarm_coordinate_function", "Function to determine particle coordinates", "DMSwarmSetCoordinateFunction", funcname, funcname, sizeof(funcname), &flg));
846d0609cedSBarry Smith   PetscOptionsEnd();
8476c5a79ebSMatthew G. Knepley   if (flg) {
8486c5a79ebSMatthew G. Knepley     PetscSimplePointFunc coordFunc;
84935b38c60SMatthew G. Knepley 
8506c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
8516c5a79ebSMatthew G. Knepley     PetscCall(PetscDLSym(NULL, funcname, (void **)&coordFunc));
8526c5a79ebSMatthew G. Knepley     PetscCheck(coordFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
8536c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
8546c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmSetLocalSizes(sw, N[rank] * Ns, 0));
8556c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmSetCoordinateFunction(sw, coordFunc));
8566c5a79ebSMatthew G. Knepley   } else {
8579566063dSJacob Faibussowitsch     PetscCall(DMGetDimension(sw, &dim));
8589566063dSJacob Faibussowitsch     PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
8599566063dSJacob Faibussowitsch     PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_coordinate_density", &pdf, NULL, NULL));
8606c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmComputeLocalSize(sw, N[rank], pdf));
8616c5a79ebSMatthew G. Knepley   }
8626c5a79ebSMatthew G. Knepley   PetscCall(PetscFree(N));
8633ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
86435b38c60SMatthew G. Knepley }
86535b38c60SMatthew G. Knepley 
86635b38c60SMatthew G. Knepley /*@
86735b38c60SMatthew G. Knepley   DMSwarmInitializeCoordinates - Determine the initial coordinates of particles for a PIC method
86835b38c60SMatthew G. Knepley 
86920f4b53cSBarry Smith   Not Collective
87035b38c60SMatthew G. Knepley 
87120f4b53cSBarry Smith   Input Parameter:
87220f4b53cSBarry Smith . sw - The `DMSWARM`
87335b38c60SMatthew G. Knepley 
87435b38c60SMatthew G. Knepley   Level: advanced
87535b38c60SMatthew G. Knepley 
87620f4b53cSBarry Smith   Note:
87720f4b53cSBarry Smith   Currently, we randomly place particles in their assigned cell
87820f4b53cSBarry Smith 
87920f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeVelocities()`
88035b38c60SMatthew G. Knepley @*/
881d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmInitializeCoordinates(DM sw)
882d71ae5a4SJacob Faibussowitsch {
8836c5a79ebSMatthew G. Knepley   PetscSimplePointFunc coordFunc;
88435b38c60SMatthew G. Knepley   PetscScalar         *weight;
8856c5a79ebSMatthew G. Knepley   PetscReal           *x;
88635b38c60SMatthew G. Knepley   PetscInt            *species;
8876c5a79ebSMatthew G. Knepley   void                *ctx;
88835b38c60SMatthew G. Knepley   PetscBool            removePoints = PETSC_TRUE;
88935b38c60SMatthew G. Knepley   PetscDataType        dtype;
8906c5a79ebSMatthew G. Knepley   PetscInt             Np, p, Ns, dim, d, bs;
89135b38c60SMatthew G. Knepley 
89235b38c60SMatthew G. Knepley   PetscFunctionBeginUser;
8936c5a79ebSMatthew G. Knepley   PetscCall(DMGetDimension(sw, &dim));
8946c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetLocalSize(sw, &Np));
8959566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
8966c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));
8976c5a79ebSMatthew G. Knepley 
8986c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, &bs, &dtype, (void **)&x));
8996c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight));
9006c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
9016c5a79ebSMatthew G. Knepley   if (coordFunc) {
9026c5a79ebSMatthew G. Knepley     PetscCall(DMGetApplicationContext(sw, &ctx));
9036c5a79ebSMatthew G. Knepley     for (p = 0; p < Np; ++p) {
9046c5a79ebSMatthew G. Knepley       PetscScalar X[3];
9056c5a79ebSMatthew G. Knepley 
9066c5a79ebSMatthew G. Knepley       PetscCall((*coordFunc)(dim, 0., NULL, p, X, ctx));
9076c5a79ebSMatthew G. Knepley       for (d = 0; d < dim; ++d) x[p * dim + d] = PetscRealPart(X[d]);
9086c5a79ebSMatthew G. Knepley       weight[p]  = 1.0;
9096c5a79ebSMatthew G. Knepley       species[p] = p % Ns;
9106c5a79ebSMatthew G. Knepley     }
9116c5a79ebSMatthew G. Knepley   } else {
9126c5a79ebSMatthew G. Knepley     DM          dm;
9136c5a79ebSMatthew G. Knepley     PetscRandom rnd;
9146c5a79ebSMatthew G. Knepley     PetscReal   xi0[3];
9156c5a79ebSMatthew G. Knepley     PetscInt    cStart, cEnd, c;
9166c5a79ebSMatthew G. Knepley 
9179566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetCellDM(sw, &dm));
9189566063dSJacob Faibussowitsch     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
919ea7032a0SMatthew G. Knepley     PetscCall(DMGetApplicationContext(sw, &ctx));
92035b38c60SMatthew G. Knepley 
92135b38c60SMatthew G. Knepley     /* Set particle position randomly in cell, set weights to 1 */
9229566063dSJacob Faibussowitsch     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd));
9239566063dSJacob Faibussowitsch     PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
9249566063dSJacob Faibussowitsch     PetscCall(PetscRandomSetFromOptions(rnd));
9259566063dSJacob Faibussowitsch     PetscCall(DMSwarmSortGetAccess(sw));
92635b38c60SMatthew G. Knepley     for (d = 0; d < dim; ++d) xi0[d] = -1.0;
92735b38c60SMatthew G. Knepley     for (c = cStart; c < cEnd; ++c) {
92835b38c60SMatthew G. Knepley       PetscReal v0[3], J[9], invJ[9], detJ;
92935b38c60SMatthew G. Knepley       PetscInt *pidx, Npc, q;
93035b38c60SMatthew G. Knepley 
9319566063dSJacob Faibussowitsch       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
9329566063dSJacob Faibussowitsch       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ));
93335b38c60SMatthew G. Knepley       for (q = 0; q < Npc; ++q) {
93435b38c60SMatthew G. Knepley         const PetscInt p = pidx[q];
93535b38c60SMatthew G. Knepley         PetscReal      xref[3];
93635b38c60SMatthew G. Knepley 
9379566063dSJacob Faibussowitsch         for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &xref[d]));
93835b38c60SMatthew G. Knepley         CoordinatesRefToReal(dim, dim, xi0, v0, J, xref, &x[p * dim]);
93935b38c60SMatthew G. Knepley 
940ea7032a0SMatthew G. Knepley         weight[p]  = 1.0 / Np;
9416c5a79ebSMatthew G. Knepley         species[p] = p % Ns;
94235b38c60SMatthew G. Knepley       }
9439566063dSJacob Faibussowitsch       PetscCall(PetscFree(pidx));
94435b38c60SMatthew G. Knepley     }
9459566063dSJacob Faibussowitsch     PetscCall(PetscRandomDestroy(&rnd));
9469566063dSJacob Faibussowitsch     PetscCall(DMSwarmSortRestoreAccess(sw));
9476c5a79ebSMatthew G. Knepley   }
9489566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
9496c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
9509566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
9516c5a79ebSMatthew G. Knepley 
9529566063dSJacob Faibussowitsch   PetscCall(DMSwarmMigrate(sw, removePoints));
9539566063dSJacob Faibussowitsch   PetscCall(DMLocalizeCoordinates(sw));
9543ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
95535b38c60SMatthew G. Knepley }
95635b38c60SMatthew G. Knepley 
95735b38c60SMatthew G. Knepley /*@C
95835b38c60SMatthew G. Knepley   DMSwarmInitializeVelocities - Set the initial velocities of particles using a distribution.
95935b38c60SMatthew G. Knepley 
96020f4b53cSBarry Smith   Collective
96135b38c60SMatthew G. Knepley 
96235b38c60SMatthew G. Knepley   Input Parameters:
96320f4b53cSBarry Smith + sw      - The `DMSWARM` object
96435b38c60SMatthew G. Knepley . sampler - A function which uniformly samples the velocity PDF
96535b38c60SMatthew G. Knepley - v0      - The velocity scale for nondimensionalization for each species
96635b38c60SMatthew G. Knepley 
96735b38c60SMatthew G. Knepley   Level: advanced
96835b38c60SMatthew G. Knepley 
96920f4b53cSBarry Smith   Note:
97020f4b53cSBarry Smith   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.
97120f4b53cSBarry Smith 
97220f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocitiesFromOptions()`
97335b38c60SMatthew G. Knepley @*/
974d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmInitializeVelocities(DM sw, PetscProbFunc sampler, const PetscReal v0[])
975d71ae5a4SJacob Faibussowitsch {
9766c5a79ebSMatthew G. Knepley   PetscSimplePointFunc velFunc;
97735b38c60SMatthew G. Knepley   PetscReal           *v;
97835b38c60SMatthew G. Knepley   PetscInt            *species;
9796c5a79ebSMatthew G. Knepley   void                *ctx;
98035b38c60SMatthew G. Knepley   PetscInt             dim, Np, p;
98135b38c60SMatthew G. Knepley 
98235b38c60SMatthew G. Knepley   PetscFunctionBegin;
9836c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));
98435b38c60SMatthew G. Knepley 
9859566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(sw, &dim));
9869566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetLocalSize(sw, &Np));
9879566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
9889566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
9896c5a79ebSMatthew G. Knepley   if (v0[0] == 0.) {
9906c5a79ebSMatthew G. Knepley     PetscCall(PetscArrayzero(v, Np * dim));
9916c5a79ebSMatthew G. Knepley   } else if (velFunc) {
9926c5a79ebSMatthew G. Knepley     PetscCall(DMGetApplicationContext(sw, &ctx));
9936c5a79ebSMatthew G. Knepley     for (p = 0; p < Np; ++p) {
9946c5a79ebSMatthew G. Knepley       PetscInt    s = species[p], d;
9956c5a79ebSMatthew G. Knepley       PetscScalar vel[3];
9966c5a79ebSMatthew G. Knepley 
9976c5a79ebSMatthew G. Knepley       PetscCall((*velFunc)(dim, 0., NULL, p, vel, ctx));
9986c5a79ebSMatthew G. Knepley       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
9996c5a79ebSMatthew G. Knepley     }
10006c5a79ebSMatthew G. Knepley   } else {
10016c5a79ebSMatthew G. Knepley     PetscRandom rnd;
10026c5a79ebSMatthew G. Knepley 
10036c5a79ebSMatthew G. Knepley     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
10046c5a79ebSMatthew G. Knepley     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
10056c5a79ebSMatthew G. Knepley     PetscCall(PetscRandomSetFromOptions(rnd));
10066c5a79ebSMatthew G. Knepley 
100735b38c60SMatthew G. Knepley     for (p = 0; p < Np; ++p) {
100835b38c60SMatthew G. Knepley       PetscInt  s = species[p], d;
100935b38c60SMatthew G. Knepley       PetscReal a[3], vel[3];
101035b38c60SMatthew G. Knepley 
10119566063dSJacob Faibussowitsch       for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &a[d]));
10129566063dSJacob Faibussowitsch       PetscCall(sampler(a, NULL, vel));
1013ad540459SPierre Jolivet       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * vel[d];
101435b38c60SMatthew G. Knepley     }
10156c5a79ebSMatthew G. Knepley     PetscCall(PetscRandomDestroy(&rnd));
10166c5a79ebSMatthew G. Knepley   }
10179566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
10189566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
10193ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
102035b38c60SMatthew G. Knepley }
102135b38c60SMatthew G. Knepley 
102235b38c60SMatthew G. Knepley /*@
102335b38c60SMatthew G. Knepley   DMSwarmInitializeVelocitiesFromOptions - Set the initial velocities of particles using a distribution determined from options.
102435b38c60SMatthew G. Knepley 
102520f4b53cSBarry Smith   Collective
102635b38c60SMatthew G. Knepley 
102735b38c60SMatthew G. Knepley   Input Parameters:
102820f4b53cSBarry Smith + sw - The `DMSWARM` object
102935b38c60SMatthew G. Knepley - v0 - The velocity scale for nondimensionalization for each species
103035b38c60SMatthew G. Knepley 
103135b38c60SMatthew G. Knepley   Level: advanced
103235b38c60SMatthew G. Knepley 
103320f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
103435b38c60SMatthew G. Knepley @*/
1035d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
1036d71ae5a4SJacob Faibussowitsch {
103735b38c60SMatthew G. Knepley   PetscProbFunc sampler;
103835b38c60SMatthew G. Knepley   PetscInt      dim;
103935b38c60SMatthew G. Knepley   const char   *prefix;
10406c5a79ebSMatthew G. Knepley   char          funcname[PETSC_MAX_PATH_LEN];
10416c5a79ebSMatthew G. Knepley   PetscBool     flg;
104235b38c60SMatthew G. Knepley 
104335b38c60SMatthew G. Knepley   PetscFunctionBegin;
1044d0609cedSBarry Smith   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
10456c5a79ebSMatthew G. Knepley   PetscCall(PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg));
1046d0609cedSBarry Smith   PetscOptionsEnd();
10476c5a79ebSMatthew G. Knepley   if (flg) {
10486c5a79ebSMatthew G. Knepley     PetscSimplePointFunc velFunc;
10496c5a79ebSMatthew G. Knepley 
10506c5a79ebSMatthew G. Knepley     PetscCall(PetscDLSym(NULL, funcname, (void **)&velFunc));
10516c5a79ebSMatthew G. Knepley     PetscCheck(velFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
10526c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmSetVelocityFunction(sw, velFunc));
10536c5a79ebSMatthew G. Knepley   }
10549566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(sw, &dim));
10559566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
10569566063dSJacob Faibussowitsch   PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_velocity_density", NULL, NULL, &sampler));
10579566063dSJacob Faibussowitsch   PetscCall(DMSwarmInitializeVelocities(sw, sampler, v0));
10583ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
105935b38c60SMatthew G. Knepley }
1060