xref: /petsc/src/dm/impls/swarm/swarmpic.c (revision 20f4b53cbb5e9bd9ef12b76a8697d60d197cda17)
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
23*20f4b53cSBarry Smith    DMSwarmSetPointsUniformCoordinates - Set point coordinates in a `DMSWARM` on a regular (ijk) grid
240e2ec84fSDave May 
25*20f4b53cSBarry Smith    Collective
260e2ec84fSDave May 
270e2ec84fSDave May    Input parameters:
28*20f4b53cSBarry 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)
32*20f4b53cSBarry 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:
37*20f4b53cSBarry Smith    When using mode = `INSERT_VALUES`, this method will reset the number of particles in the `DMSWARM`
38*20f4b53cSBarry Smith    to be npoints[0]*npoints[1] (2D) or npoints[0]*npoints[1]*npoints[2] (3D). When using mode = `ADD_VALUES`,
39*20f4b53cSBarry Smith    new points will be appended to any already existing in the `DMSWARM`
400e2ec84fSDave May 
41*20f4b53cSBarry 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
181*20f4b53cSBarry Smith    DMSwarmSetPointCoordinates - Set point coordinates in a `DMSWARM` from a user defined list
1820e2ec84fSDave May 
183*20f4b53cSBarry Smith    Collective
1840e2ec84fSDave May 
1850e2ec84fSDave May    Input parameters:
186*20f4b53cSBarry Smith +  dm - the `DMSWARM`
1870e2ec84fSDave May .  npoints - the number of points to insert
1880e2ec84fSDave May .  coor - the coordinate values
189*20f4b53cSBarry 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
190*20f4b53cSBarry 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:
195*20f4b53cSBarry Smith    If the user has specified `redundant` as `PETSC_FALSE`, the cell `DM` will attempt to locate the coordinates provided by `coor` within
196*20f4b53cSBarry 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
197*20f4b53cSBarry Smith    added to the `DMSWARM`.
1980e2ec84fSDave May 
199*20f4b53cSBarry 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 
340*20f4b53cSBarry Smith    Not Collective
3410e2ec84fSDave May 
3420e2ec84fSDave May    Input parameters:
343*20f4b53cSBarry Smith +  dm - the `DMSWARM`
344*20f4b53cSBarry 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:
350*20f4b53cSBarry Smith    The insert method will reset any previous defined points within the `DMSWARM`.
351e7af74c8SDave May 
352*20f4b53cSBarry 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*20f4b53cSBarry Smith    When using a `DMPLEX` the following case are supported:
355*20f4b53cSBarry 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.
359*20f4b53cSBarry Smith .ve
3600e2ec84fSDave May 
361*20f4b53cSBarry 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 
386*20f4b53cSBarry Smith    Not Collective
387431382f2SDave May 
388431382f2SDave May    Input parameters:
389*20f4b53cSBarry Smith +  dm - the `DMSWARM`
390*20f4b53cSBarry Smith .  celldm - the cell `DM`
391431382f2SDave May .  npoints - the number of points to insert in each cell
392431382f2SDave May -  xi - the coordinates (defined in the local coordinate system for each cell) to insert
393431382f2SDave May 
394431382f2SDave May  Level: beginner
395431382f2SDave May 
396431382f2SDave May  Notes:
397*20f4b53cSBarry Smith  The method will reset any previous defined points within the `DMSWARM`.
398*20f4b53cSBarry Smith  Only supported for `DMPLEX`. If you are using a `DMDA` it is recommended to either use
399*20f4b53cSBarry Smith  `DMSwarmInsertPointsUsingCellDM()`, or extract and set the coordinates yourself the following code
400*20f4b53cSBarry Smith .vb
401*20f4b53cSBarry Smith     PetscReal *coor;
402*20f4b53cSBarry Smith     DMSwarmGetField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
403*20f4b53cSBarry Smith     // user code to define the coordinates here
404*20f4b53cSBarry Smith     DMSwarmRestoreField(dm,DMSwarmPICField_coor,NULL,NULL,(void**)&coor);
405*20f4b53cSBarry Smith .ve
406e7af74c8SDave May 
407*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetCellDM()`, `DMSwarmInsertPointsUsingCellDM()`
408431382f2SDave May @*/
409d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmSetPointCoordinatesCellwise(DM dm, PetscInt npoints, PetscReal xi[])
410d71ae5a4SJacob Faibussowitsch {
411431382f2SDave May   DM        celldm;
412431382f2SDave May   PetscBool isDA, isPLEX;
413431382f2SDave May 
4140e2ec84fSDave May   PetscFunctionBegin;
415431382f2SDave May   DMSWARMPICVALID(dm);
4169566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dm, &celldm));
4179566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
4189566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
41928b400f6SJacob Faibussowitsch   PetscCheck(!isDA, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMPLEX. Recommended you use DMSwarmInsertPointsUsingCellDM()");
420f7d195e4SLawrence Mitchell   if (isPLEX) {
4219566063dSJacob Faibussowitsch     PetscCall(private_DMSwarmSetPointCoordinatesCellwise_PLEX(dm, celldm, npoints, xi));
422431382f2SDave May   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
4233ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
4240e2ec84fSDave May }
425431382f2SDave May 
4260e2ec84fSDave May /* Field projection API */
42777048351SPatrick Sanan extern PetscErrorCode private_DMSwarmProjectFields_DA(DM swarm, DM celldm, PetscInt project_type, PetscInt nfields, DMSwarmDataField dfield[], Vec vecs[]);
42877048351SPatrick Sanan extern PetscErrorCode private_DMSwarmProjectFields_PLEX(DM swarm, DM celldm, PetscInt project_type, PetscInt nfields, DMSwarmDataField dfield[], Vec vecs[]);
42994f7d2dcSDave May 
43094f7d2dcSDave May /*@C
431*20f4b53cSBarry Smith    DMSwarmProjectFields - Project a set of swarm fields onto the cell `DM`
43294f7d2dcSDave May 
433*20f4b53cSBarry Smith    Collective
43494f7d2dcSDave May 
43594f7d2dcSDave May    Input parameters:
436*20f4b53cSBarry Smith +  dm - the `DMSWARM`
43794f7d2dcSDave May .  nfields - the number of swarm fields to project
43894f7d2dcSDave May .  fieldnames - the textual names of the swarm fields to project
43994f7d2dcSDave May .  fields - an array of Vec's of length nfields
44094f7d2dcSDave May -  reuse - flag indicating whether the array and contents of fields should be re-used or internally allocated
44194f7d2dcSDave May 
442*20f4b53cSBarry Smith    Level: beginner
443*20f4b53cSBarry Smith 
444*20f4b53cSBarry Smith    Notes:
44594f7d2dcSDave May    Currently, the only available projection method consists of
446*20f4b53cSBarry Smith .vb
44794f7d2dcSDave May      phi_i = \sum_{p=0}^{np} N_i(x_p) phi_p dJ / \sum_{p=0}^{np} N_i(x_p) dJ
44894f7d2dcSDave May    where phi_p is the swarm field at point p,
44994f7d2dcSDave May      N_i() is the cell DM basis function at vertex i,
45094f7d2dcSDave May      dJ is the determinant of the cell Jacobian and
45194f7d2dcSDave May      phi_i is the projected vertex value of the field phi.
452*20f4b53cSBarry Smith .ve
45394f7d2dcSDave May 
454*20f4b53cSBarry Smith    If `reuse` is `PETSC_FALSE`, this function will allocate the array of `Vec`'s, and each individual `Vec`.
455*20f4b53cSBarry Smith      The user is responsible for destroying both the array and the individual `Vec` objects.
45694f7d2dcSDave May 
457*20f4b53cSBarry Smith    Only swarm fields registered with data type of `PETSC_REAL` can be projected onto the cell `DM`.
458e7af74c8SDave May 
459e7af74c8SDave May    Only swarm fields of block size = 1 can currently be projected.
460e7af74c8SDave May 
461*20f4b53cSBarry Smith    The only projection methods currently only support the `DMDA` (2D) and `DMPLEX` (triangles 2D).
46294f7d2dcSDave May 
463*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
46494f7d2dcSDave May @*/
465d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmProjectFields(DM dm, PetscInt nfields, const char *fieldnames[], Vec **fields, PetscBool reuse)
466d71ae5a4SJacob Faibussowitsch {
46794f7d2dcSDave May   DM_Swarm         *swarm = (DM_Swarm *)dm->data;
46877048351SPatrick Sanan   DMSwarmDataField *gfield;
46994f7d2dcSDave May   DM                celldm;
47094f7d2dcSDave May   PetscBool         isDA, isPLEX;
47194f7d2dcSDave May   Vec              *vecs;
47294f7d2dcSDave May   PetscInt          f, nvecs;
47394f7d2dcSDave May   PetscInt          project_type = 0;
47494f7d2dcSDave May 
4750e2ec84fSDave May   PetscFunctionBegin;
47694f7d2dcSDave May   DMSWARMPICVALID(dm);
4779566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(dm, &celldm));
4789566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(nfields, &gfield));
47994f7d2dcSDave May   nvecs = 0;
48094f7d2dcSDave May   for (f = 0; f < nfields; f++) {
4819566063dSJacob Faibussowitsch     PetscCall(DMSwarmDataBucketGetDMSwarmDataFieldByName(swarm->db, fieldnames[f], &gfield[f]));
4825f80ce2aSJacob 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");
4835f80ce2aSJacob Faibussowitsch     PetscCheck(gfield[f]->bs == 1, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Projection only valid for fields with block size = 1");
48494f7d2dcSDave May     nvecs += gfield[f]->bs;
48594f7d2dcSDave May   }
48694f7d2dcSDave May   if (!reuse) {
4879566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(nvecs, &vecs));
48894f7d2dcSDave May     for (f = 0; f < nvecs; f++) {
4899566063dSJacob Faibussowitsch       PetscCall(DMCreateGlobalVector(celldm, &vecs[f]));
4909566063dSJacob Faibussowitsch       PetscCall(PetscObjectSetName((PetscObject)vecs[f], gfield[f]->name));
49194f7d2dcSDave May     }
49294f7d2dcSDave May   } else {
49394f7d2dcSDave May     vecs = *fields;
49494f7d2dcSDave May   }
49594f7d2dcSDave May 
4969566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isDA));
4979566063dSJacob Faibussowitsch   PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isPLEX));
49894f7d2dcSDave May   if (isDA) {
4999566063dSJacob Faibussowitsch     PetscCall(private_DMSwarmProjectFields_DA(dm, celldm, project_type, nfields, gfield, vecs));
50094f7d2dcSDave May   } else if (isPLEX) {
5019566063dSJacob Faibussowitsch     PetscCall(private_DMSwarmProjectFields_PLEX(dm, celldm, project_type, nfields, gfield, vecs));
50294f7d2dcSDave May   } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Only supported for cell DMs of type DMDA and DMPLEX");
50394f7d2dcSDave May 
5049566063dSJacob Faibussowitsch   PetscCall(PetscFree(gfield));
5055f80ce2aSJacob Faibussowitsch   if (!reuse) *fields = vecs;
5063ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5070e2ec84fSDave May }
5080e2ec84fSDave May 
5090e2ec84fSDave May /*@C
510b799feefSDave May    DMSwarmCreatePointPerCellCount - Count the number of points within all cells in the cell DM
5110e2ec84fSDave May 
512*20f4b53cSBarry Smith    Not Collective
5130e2ec84fSDave May 
5140e2ec84fSDave May    Input parameter:
515*20f4b53cSBarry Smith .  dm - the `DMSWARM`
5160e2ec84fSDave May 
5170e2ec84fSDave May    Output parameters:
518*20f4b53cSBarry Smith +  ncells - the number of cells in the cell `DM` (optional argument, pass `NULL` to ignore)
519b799feefSDave May -  count - array of length ncells containing the number of points per cell
5200e2ec84fSDave May 
5210e2ec84fSDave May    Level: beginner
5220e2ec84fSDave May 
5230e2ec84fSDave May    Notes:
5240e2ec84fSDave May    The array count is allocated internally and must be free'd by the user.
5250e2ec84fSDave May 
526*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetType()`, `DMSwarmSetCellDM()`, `DMSwarmType`
5270e2ec84fSDave May @*/
528d71ae5a4SJacob Faibussowitsch PETSC_EXTERN PetscErrorCode DMSwarmCreatePointPerCellCount(DM dm, PetscInt *ncells, PetscInt **count)
529d71ae5a4SJacob Faibussowitsch {
530b799feefSDave May   PetscBool isvalid;
531b799feefSDave May   PetscInt  nel;
532b799feefSDave May   PetscInt *sum;
533b799feefSDave May 
5340e2ec84fSDave May   PetscFunctionBegin;
5359566063dSJacob Faibussowitsch   PetscCall(DMSwarmSortGetIsValid(dm, &isvalid));
536b799feefSDave May   nel = 0;
537b799feefSDave May   if (isvalid) {
538b799feefSDave May     PetscInt e;
539b799feefSDave May 
5409566063dSJacob Faibussowitsch     PetscCall(DMSwarmSortGetSizes(dm, &nel, NULL));
541b799feefSDave May 
5429566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(nel, &sum));
54348a46eb9SPierre Jolivet     for (e = 0; e < nel; e++) PetscCall(DMSwarmSortGetNumberOfPointsPerCell(dm, e, &sum[e]));
544b799feefSDave May   } else {
545b799feefSDave May     DM        celldm;
546b799feefSDave May     PetscBool isda, isplex, isshell;
547b799feefSDave May     PetscInt  p, npoints;
548b799feefSDave May     PetscInt *swarm_cellid;
549b799feefSDave May 
550b799feefSDave May     /* get the number of cells */
5519566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetCellDM(dm, &celldm));
5529566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMDA, &isda));
5539566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMPLEX, &isplex));
5549566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)celldm, DMSHELL, &isshell));
555b799feefSDave May     if (isda) {
556b799feefSDave May       PetscInt        _nel, _npe;
557b799feefSDave May       const PetscInt *_element;
558b799feefSDave May 
5599566063dSJacob Faibussowitsch       PetscCall(DMDAGetElements(celldm, &_nel, &_npe, &_element));
560b799feefSDave May       nel = _nel;
5619566063dSJacob Faibussowitsch       PetscCall(DMDARestoreElements(celldm, &_nel, &_npe, &_element));
562b799feefSDave May     } else if (isplex) {
563b799feefSDave May       PetscInt ps, pe;
564b799feefSDave May 
5659566063dSJacob Faibussowitsch       PetscCall(DMPlexGetHeightStratum(celldm, 0, &ps, &pe));
566b799feefSDave May       nel = pe - ps;
567b799feefSDave May     } else if (isshell) {
568b799feefSDave May       PetscErrorCode (*method_DMShellGetNumberOfCells)(DM, PetscInt *);
569b799feefSDave May 
5709566063dSJacob Faibussowitsch       PetscCall(PetscObjectQueryFunction((PetscObject)celldm, "DMGetNumberOfCells_C", &method_DMShellGetNumberOfCells));
571b799feefSDave May       if (method_DMShellGetNumberOfCells) {
5729566063dSJacob Faibussowitsch         PetscCall(method_DMShellGetNumberOfCells(celldm, &nel));
5739371c9d4SSatish Balay       } else
5749371c9d4SSatish 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);");
575b799feefSDave 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");
576b799feefSDave May 
5779566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(nel, &sum));
5789566063dSJacob Faibussowitsch     PetscCall(PetscArrayzero(sum, nel));
5799566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetLocalSize(dm, &npoints));
5809566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
581b799feefSDave May     for (p = 0; p < npoints; p++) {
582ad540459SPierre Jolivet       if (swarm_cellid[p] != DMLOCATEPOINT_POINT_NOT_FOUND) sum[swarm_cellid[p]]++;
583b799feefSDave May     }
5849566063dSJacob Faibussowitsch     PetscCall(DMSwarmRestoreField(dm, DMSwarmPICField_cellid, NULL, NULL, (void **)&swarm_cellid));
585b799feefSDave May   }
586ad540459SPierre Jolivet   if (ncells) *ncells = nel;
587b799feefSDave May   *count = sum;
5883ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5890e2ec84fSDave May }
59035b38c60SMatthew G. Knepley 
59135b38c60SMatthew G. Knepley /*@
59235b38c60SMatthew G. Knepley   DMSwarmGetNumSpecies - Get the number of particle species
59335b38c60SMatthew G. Knepley 
594*20f4b53cSBarry Smith   Not Collective
59535b38c60SMatthew G. Knepley 
59635b38c60SMatthew G. Knepley   Input parameter:
597*20f4b53cSBarry Smith . dm - the `DMSWARM`
59835b38c60SMatthew G. Knepley 
59935b38c60SMatthew G. Knepley   Output parameters:
60035b38c60SMatthew G. Knepley . Ns - the number of species
60135b38c60SMatthew G. Knepley 
60235b38c60SMatthew G. Knepley   Level: intermediate
60335b38c60SMatthew G. Knepley 
604*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
60535b38c60SMatthew G. Knepley @*/
606d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmGetNumSpecies(DM sw, PetscInt *Ns)
607d71ae5a4SJacob Faibussowitsch {
60835b38c60SMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
60935b38c60SMatthew G. Knepley 
61035b38c60SMatthew G. Knepley   PetscFunctionBegin;
61135b38c60SMatthew G. Knepley   *Ns = swarm->Ns;
6123ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
61335b38c60SMatthew G. Knepley }
61435b38c60SMatthew G. Knepley 
61535b38c60SMatthew G. Knepley /*@
61635b38c60SMatthew G. Knepley   DMSwarmSetNumSpecies - Set the number of particle species
61735b38c60SMatthew G. Knepley 
618*20f4b53cSBarry Smith   Not Collective
61935b38c60SMatthew G. Knepley 
62035b38c60SMatthew G. Knepley   Input parameter:
621*20f4b53cSBarry Smith + dm - the `DMSWARM`
62235b38c60SMatthew G. Knepley - Ns - the number of species
62335b38c60SMatthew G. Knepley 
62435b38c60SMatthew G. Knepley   Level: intermediate
62535b38c60SMatthew G. Knepley 
626*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmGetNumSpecies()`, `DMSwarmSetType()`, `DMSwarmType`
62735b38c60SMatthew G. Knepley @*/
628d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmSetNumSpecies(DM sw, PetscInt Ns)
629d71ae5a4SJacob Faibussowitsch {
63035b38c60SMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
63135b38c60SMatthew G. Knepley 
63235b38c60SMatthew G. Knepley   PetscFunctionBegin;
63335b38c60SMatthew G. Knepley   swarm->Ns = Ns;
6343ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
63535b38c60SMatthew G. Knepley }
63635b38c60SMatthew G. Knepley 
63735b38c60SMatthew G. Knepley /*@C
6386c5a79ebSMatthew G. Knepley   DMSwarmGetCoordinateFunction - Get the function setting initial particle positions, if it exists
6396c5a79ebSMatthew G. Knepley 
640*20f4b53cSBarry Smith   Not Collective
6416c5a79ebSMatthew G. Knepley 
6426c5a79ebSMatthew G. Knepley   Input parameter:
643*20f4b53cSBarry Smith . dm - the `DMSWARM`
6446c5a79ebSMatthew G. Knepley 
6456c5a79ebSMatthew G. Knepley   Output Parameter:
646*20f4b53cSBarry Smith . coordFunc - the function setting initial particle positions, or `NULL`
6476c5a79ebSMatthew G. Knepley 
6486c5a79ebSMatthew G. Knepley   Level: intermediate
6496c5a79ebSMatthew G. Knepley 
650*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetCoordinateFunction()`, `DMSwarmGetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
6516c5a79ebSMatthew G. Knepley @*/
652d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmGetCoordinateFunction(DM sw, PetscSimplePointFunc *coordFunc)
653d71ae5a4SJacob Faibussowitsch {
6546c5a79ebSMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
6556c5a79ebSMatthew G. Knepley 
6566c5a79ebSMatthew G. Knepley   PetscFunctionBegin;
6576c5a79ebSMatthew G. Knepley   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
6586c5a79ebSMatthew G. Knepley   PetscValidPointer(coordFunc, 2);
6596c5a79ebSMatthew G. Knepley   *coordFunc = swarm->coordFunc;
6603ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
6616c5a79ebSMatthew G. Knepley }
6626c5a79ebSMatthew G. Knepley 
6636c5a79ebSMatthew G. Knepley /*@C
6646c5a79ebSMatthew G. Knepley   DMSwarmSetCoordinateFunction - Set the function setting initial particle positions
6656c5a79ebSMatthew G. Knepley 
666*20f4b53cSBarry Smith   Not Collective
6676c5a79ebSMatthew G. Knepley 
6686c5a79ebSMatthew G. Knepley   Input parameters:
669*20f4b53cSBarry Smith + dm - the `DMSWARM`
6706c5a79ebSMatthew G. Knepley - coordFunc - the function setting initial particle positions
6716c5a79ebSMatthew G. Knepley 
6726c5a79ebSMatthew G. Knepley   Level: intermediate
6736c5a79ebSMatthew G. Knepley 
674*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmGetCoordinateFunction()`, `DMSwarmSetVelocityFunction()`, `DMSwarmInitializeCoordinates()`
6756c5a79ebSMatthew G. Knepley @*/
676d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmSetCoordinateFunction(DM sw, PetscSimplePointFunc coordFunc)
677d71ae5a4SJacob Faibussowitsch {
6786c5a79ebSMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
6796c5a79ebSMatthew G. Knepley 
6806c5a79ebSMatthew G. Knepley   PetscFunctionBegin;
6816c5a79ebSMatthew G. Knepley   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
6826c5a79ebSMatthew G. Knepley   PetscValidFunction(coordFunc, 2);
6836c5a79ebSMatthew G. Knepley   swarm->coordFunc = coordFunc;
6843ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
6856c5a79ebSMatthew G. Knepley }
6866c5a79ebSMatthew G. Knepley 
6876c5a79ebSMatthew G. Knepley /*@C
6886c5a79ebSMatthew G. Knepley   DMSwarmGetCoordinateFunction - Get the function setting initial particle velocities, if it exists
6896c5a79ebSMatthew G. Knepley 
690*20f4b53cSBarry Smith   Not Collective
6916c5a79ebSMatthew G. Knepley 
6926c5a79ebSMatthew G. Knepley   Input parameter:
693*20f4b53cSBarry Smith . dm - the `DMSWARM`
6946c5a79ebSMatthew G. Knepley 
6956c5a79ebSMatthew G. Knepley   Output Parameter:
696*20f4b53cSBarry Smith . velFunc - the function setting initial particle velocities, or `NULL`
6976c5a79ebSMatthew G. Knepley 
6986c5a79ebSMatthew G. Knepley   Level: intermediate
6996c5a79ebSMatthew G. Knepley 
700*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmSetVelocityFunction()`, `DMSwarmGetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
7016c5a79ebSMatthew G. Knepley @*/
702d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmGetVelocityFunction(DM sw, PetscSimplePointFunc *velFunc)
703d71ae5a4SJacob Faibussowitsch {
7046c5a79ebSMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
7056c5a79ebSMatthew G. Knepley 
7066c5a79ebSMatthew G. Knepley   PetscFunctionBegin;
7076c5a79ebSMatthew G. Knepley   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
7086c5a79ebSMatthew G. Knepley   PetscValidPointer(velFunc, 2);
7096c5a79ebSMatthew G. Knepley   *velFunc = swarm->velFunc;
7103ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7116c5a79ebSMatthew G. Knepley }
7126c5a79ebSMatthew G. Knepley 
7136c5a79ebSMatthew G. Knepley /*@C
7146c5a79ebSMatthew G. Knepley   DMSwarmSetVelocityFunction - Set the function setting initial particle velocities
7156c5a79ebSMatthew G. Knepley 
716*20f4b53cSBarry Smith   Not Collective
7176c5a79ebSMatthew G. Knepley 
7186c5a79ebSMatthew G. Knepley   Input parameters:
719*20f4b53cSBarry Smith + dm - the `DMSWARM`
7206c5a79ebSMatthew G. Knepley - coordFunc - the function setting initial particle velocities
7216c5a79ebSMatthew G. Knepley 
7226c5a79ebSMatthew G. Knepley   Level: intermediate
7236c5a79ebSMatthew G. Knepley 
724*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmGetVelocityFunction()`, `DMSwarmSetCoordinateFunction()`, `DMSwarmInitializeVelocities()`
7256c5a79ebSMatthew G. Knepley @*/
726d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmSetVelocityFunction(DM sw, PetscSimplePointFunc velFunc)
727d71ae5a4SJacob Faibussowitsch {
7286c5a79ebSMatthew G. Knepley   DM_Swarm *swarm = (DM_Swarm *)sw->data;
7296c5a79ebSMatthew G. Knepley 
7306c5a79ebSMatthew G. Knepley   PetscFunctionBegin;
7316c5a79ebSMatthew G. Knepley   PetscValidHeaderSpecific(sw, DM_CLASSID, 1);
7326c5a79ebSMatthew G. Knepley   PetscValidFunction(velFunc, 2);
7336c5a79ebSMatthew G. Knepley   swarm->velFunc = velFunc;
7343ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
7356c5a79ebSMatthew G. Knepley }
7366c5a79ebSMatthew G. Knepley 
7376c5a79ebSMatthew G. Knepley /*@C
73835b38c60SMatthew G. Knepley   DMSwarmComputeLocalSize - Compute the local number and distribution of particles based upon a density function
73935b38c60SMatthew G. Knepley 
740*20f4b53cSBarry Smith   Not Collective
74135b38c60SMatthew G. Knepley 
74235b38c60SMatthew G. Knepley   Input Parameters:
743*20f4b53cSBarry Smith + sw      - The `DMSWARM`
74435b38c60SMatthew G. Knepley . N       - The target number of particles
74535b38c60SMatthew G. Knepley - density - The density field for the particle layout, normalized to unity
74635b38c60SMatthew G. Knepley 
74735b38c60SMatthew G. Knepley   Level: advanced
74835b38c60SMatthew G. Knepley 
749*20f4b53cSBarry Smith   Note:
750*20f4b53cSBarry Smith   One particle will be created for each species.
751*20f4b53cSBarry Smith 
752*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSizeFromOptions()`
75335b38c60SMatthew G. Knepley @*/
754d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmComputeLocalSize(DM sw, PetscInt N, PetscProbFunc density)
755d71ae5a4SJacob Faibussowitsch {
75635b38c60SMatthew G. Knepley   DM               dm;
75735b38c60SMatthew G. Knepley   PetscQuadrature  quad;
75835b38c60SMatthew G. Knepley   const PetscReal *xq, *wq;
759ea7032a0SMatthew G. Knepley   PetscReal       *n_int;
760ea7032a0SMatthew G. Knepley   PetscInt        *npc_s, *cellid, Ni;
761ea7032a0SMatthew G. Knepley   PetscReal        gmin[3], gmax[3], xi0[3];
762ea7032a0SMatthew G. Knepley   PetscInt         Ns, cStart, cEnd, c, dim, d, Nq, q, Np = 0, p, s;
76335b38c60SMatthew G. Knepley   PetscBool        simplex;
76435b38c60SMatthew G. Knepley 
76535b38c60SMatthew G. Knepley   PetscFunctionBegin;
7669566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
7679566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetCellDM(sw, &dm));
7689566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(dm, &dim));
769ea7032a0SMatthew G. Knepley   PetscCall(DMGetBoundingBox(dm, gmin, gmax));
7709566063dSJacob Faibussowitsch   PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
7719566063dSJacob Faibussowitsch   PetscCall(DMPlexIsSimplex(dm, &simplex));
7726858538eSMatthew G. Knepley   PetscCall(DMGetCoordinatesLocalSetUp(dm));
7739566063dSJacob Faibussowitsch   if (simplex) PetscCall(PetscDTStroudConicalQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
7749566063dSJacob Faibussowitsch   else PetscCall(PetscDTGaussTensorQuadrature(dim, 1, 5, -1.0, 1.0, &quad));
7759566063dSJacob Faibussowitsch   PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, &xq, &wq));
776ea7032a0SMatthew G. Knepley   PetscCall(PetscCalloc2(Ns, &n_int, (cEnd - cStart) * Ns, &npc_s));
77735b38c60SMatthew G. Knepley   /* Integrate the density function to get the number of particles in each cell */
77835b38c60SMatthew G. Knepley   for (d = 0; d < dim; ++d) xi0[d] = -1.0;
77935b38c60SMatthew G. Knepley   for (c = 0; c < cEnd - cStart; ++c) {
78035b38c60SMatthew G. Knepley     const PetscInt cell = c + cStart;
781ea7032a0SMatthew G. Knepley     PetscReal      v0[3], J[9], invJ[9], detJ, detJp = 2. / (gmax[0] - gmin[0]), xr[3], den;
78235b38c60SMatthew G. Knepley 
783ea7032a0SMatthew G. Knepley     /*Have to transform quadrature points/weights to cell domain*/
7849566063dSJacob Faibussowitsch     PetscCall(DMPlexComputeCellGeometryFEM(dm, cell, NULL, v0, J, invJ, &detJ));
785ea7032a0SMatthew G. Knepley     PetscCall(PetscArrayzero(n_int, Ns));
78635b38c60SMatthew G. Knepley     for (q = 0; q < Nq; ++q) {
78735b38c60SMatthew G. Knepley       CoordinatesRefToReal(dim, dim, xi0, v0, J, &xq[q * dim], xr);
788ea7032a0SMatthew G. Knepley       /* Have to transform mesh to domain of definition of PDF, [-1, 1], and weight PDF by |J|/2 */
789ea7032a0SMatthew G. Knepley       xr[0] = detJp * (xr[0] - gmin[0]) - 1.;
790ea7032a0SMatthew G. Knepley 
791ea7032a0SMatthew G. Knepley       for (s = 0; s < Ns; ++s) {
792ea7032a0SMatthew G. Knepley         PetscCall(density(xr, NULL, &den));
793ea7032a0SMatthew G. Knepley         n_int[s] += (detJp * den) * (detJ * wq[q]) / (PetscReal)Ns;
79435b38c60SMatthew G. Knepley       }
795ea7032a0SMatthew G. Knepley     }
796ea7032a0SMatthew G. Knepley     for (s = 0; s < Ns; ++s) {
797ea7032a0SMatthew G. Knepley       Ni = N;
798ea7032a0SMatthew G. Knepley       npc_s[c * Ns + s] += (PetscInt)(Ni * n_int[s]);
799ea7032a0SMatthew G. Knepley       Np += npc_s[c * Ns + s];
800ea7032a0SMatthew G. Knepley     }
80135b38c60SMatthew G. Knepley   }
8029566063dSJacob Faibussowitsch   PetscCall(PetscQuadratureDestroy(&quad));
8039566063dSJacob Faibussowitsch   PetscCall(DMSwarmSetLocalSizes(sw, Np, 0));
8049566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
80535b38c60SMatthew G. Knepley   for (c = 0, p = 0; c < cEnd - cStart; ++c) {
806ea7032a0SMatthew G. Knepley     for (s = 0; s < Ns; ++s) {
807ea7032a0SMatthew G. Knepley       for (q = 0; q < npc_s[c * Ns + s]; ++q, ++p) cellid[p] = c;
808ea7032a0SMatthew G. Knepley     }
80935b38c60SMatthew G. Knepley   }
8109566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_cellid, NULL, NULL, (void **)&cellid));
811ea7032a0SMatthew G. Knepley   PetscCall(PetscFree2(n_int, npc_s));
8123ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
81335b38c60SMatthew G. Knepley }
81435b38c60SMatthew G. Knepley 
81535b38c60SMatthew G. Knepley /*@
81635b38c60SMatthew G. Knepley   DMSwarmComputeLocalSizeFromOptions - Compute the local number and distribution of particles based upon a density function determined by options
81735b38c60SMatthew G. Knepley 
818*20f4b53cSBarry Smith   Not Collective
81935b38c60SMatthew G. Knepley 
82035b38c60SMatthew G. Knepley   Input Parameters:
821*20f4b53cSBarry Smith , sw - The `DMSWARM`
82235b38c60SMatthew G. Knepley 
82335b38c60SMatthew G. Knepley   Level: advanced
82435b38c60SMatthew G. Knepley 
825*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`
82635b38c60SMatthew G. Knepley @*/
827d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmComputeLocalSizeFromOptions(DM sw)
828d71ae5a4SJacob Faibussowitsch {
82935b38c60SMatthew G. Knepley   PetscProbFunc pdf;
83035b38c60SMatthew G. Knepley   const char   *prefix;
8316c5a79ebSMatthew G. Knepley   char          funcname[PETSC_MAX_PATH_LEN];
8326c5a79ebSMatthew G. Knepley   PetscInt     *N, Ns, dim, n;
8336c5a79ebSMatthew G. Knepley   PetscBool     flg;
8346c5a79ebSMatthew G. Knepley   PetscMPIInt   size, rank;
83535b38c60SMatthew G. Knepley 
83635b38c60SMatthew G. Knepley   PetscFunctionBegin;
8376c5a79ebSMatthew G. Knepley   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)sw), &size));
8386c5a79ebSMatthew G. Knepley   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)sw), &rank));
8396c5a79ebSMatthew G. Knepley   PetscCall(PetscCalloc1(size, &N));
840d0609cedSBarry Smith   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
8416c5a79ebSMatthew G. Knepley   n = size;
8426c5a79ebSMatthew G. Knepley   PetscCall(PetscOptionsIntArray("-dm_swarm_num_particles", "The target number of particles", "", N, &n, NULL));
8436c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
8449566063dSJacob Faibussowitsch   PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg));
8459566063dSJacob Faibussowitsch   if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns));
8466c5a79ebSMatthew G. Knepley   PetscCall(PetscOptionsString("-dm_swarm_coordinate_function", "Function to determine particle coordinates", "DMSwarmSetCoordinateFunction", funcname, funcname, sizeof(funcname), &flg));
847d0609cedSBarry Smith   PetscOptionsEnd();
8486c5a79ebSMatthew G. Knepley   if (flg) {
8496c5a79ebSMatthew G. Knepley     PetscSimplePointFunc coordFunc;
85035b38c60SMatthew G. Knepley 
8516c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
8526c5a79ebSMatthew G. Knepley     PetscCall(PetscDLSym(NULL, funcname, (void **)&coordFunc));
8536c5a79ebSMatthew G. Knepley     PetscCheck(coordFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
8546c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
8556c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmSetLocalSizes(sw, N[rank] * Ns, 0));
8566c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmSetCoordinateFunction(sw, coordFunc));
8576c5a79ebSMatthew G. Knepley   } else {
8589566063dSJacob Faibussowitsch     PetscCall(DMGetDimension(sw, &dim));
8599566063dSJacob Faibussowitsch     PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
8609566063dSJacob Faibussowitsch     PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_coordinate_density", &pdf, NULL, NULL));
8616c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmComputeLocalSize(sw, N[rank], pdf));
8626c5a79ebSMatthew G. Knepley   }
8636c5a79ebSMatthew G. Knepley   PetscCall(PetscFree(N));
8643ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
86535b38c60SMatthew G. Knepley }
86635b38c60SMatthew G. Knepley 
86735b38c60SMatthew G. Knepley /*@
86835b38c60SMatthew G. Knepley   DMSwarmInitializeCoordinates - Determine the initial coordinates of particles for a PIC method
86935b38c60SMatthew G. Knepley 
870*20f4b53cSBarry Smith   Not Collective
87135b38c60SMatthew G. Knepley 
872*20f4b53cSBarry Smith   Input Parameter:
873*20f4b53cSBarry Smith . sw - The `DMSWARM`
87435b38c60SMatthew G. Knepley 
87535b38c60SMatthew G. Knepley   Level: advanced
87635b38c60SMatthew G. Knepley 
877*20f4b53cSBarry Smith   Note:
878*20f4b53cSBarry Smith   Currently, we randomly place particles in their assigned cell
879*20f4b53cSBarry Smith 
880*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeVelocities()`
88135b38c60SMatthew G. Knepley @*/
882d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmInitializeCoordinates(DM sw)
883d71ae5a4SJacob Faibussowitsch {
8846c5a79ebSMatthew G. Knepley   PetscSimplePointFunc coordFunc;
88535b38c60SMatthew G. Knepley   PetscScalar         *weight;
8866c5a79ebSMatthew G. Knepley   PetscReal           *x;
88735b38c60SMatthew G. Knepley   PetscInt            *species;
8886c5a79ebSMatthew G. Knepley   void                *ctx;
88935b38c60SMatthew G. Knepley   PetscBool            removePoints = PETSC_TRUE;
89035b38c60SMatthew G. Knepley   PetscDataType        dtype;
8916c5a79ebSMatthew G. Knepley   PetscInt             Np, p, Ns, dim, d, bs;
89235b38c60SMatthew G. Knepley 
89335b38c60SMatthew G. Knepley   PetscFunctionBeginUser;
8946c5a79ebSMatthew G. Knepley   PetscCall(DMGetDimension(sw, &dim));
8956c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetLocalSize(sw, &Np));
8969566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetNumSpecies(sw, &Ns));
8976c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetCoordinateFunction(sw, &coordFunc));
8986c5a79ebSMatthew G. Knepley 
8996c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, &bs, &dtype, (void **)&x));
9006c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&weight));
9016c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
9026c5a79ebSMatthew G. Knepley   if (coordFunc) {
9036c5a79ebSMatthew G. Knepley     PetscCall(DMGetApplicationContext(sw, &ctx));
9046c5a79ebSMatthew G. Knepley     for (p = 0; p < Np; ++p) {
9056c5a79ebSMatthew G. Knepley       PetscScalar X[3];
9066c5a79ebSMatthew G. Knepley 
9076c5a79ebSMatthew G. Knepley       PetscCall((*coordFunc)(dim, 0., NULL, p, X, ctx));
9086c5a79ebSMatthew G. Knepley       for (d = 0; d < dim; ++d) x[p * dim + d] = PetscRealPart(X[d]);
9096c5a79ebSMatthew G. Knepley       weight[p]  = 1.0;
9106c5a79ebSMatthew G. Knepley       species[p] = p % Ns;
9116c5a79ebSMatthew G. Knepley     }
9126c5a79ebSMatthew G. Knepley   } else {
9136c5a79ebSMatthew G. Knepley     DM          dm;
9146c5a79ebSMatthew G. Knepley     PetscRandom rnd;
9156c5a79ebSMatthew G. Knepley     PetscReal   xi0[3];
9166c5a79ebSMatthew G. Knepley     PetscInt    cStart, cEnd, c;
9176c5a79ebSMatthew G. Knepley 
9189566063dSJacob Faibussowitsch     PetscCall(DMSwarmGetCellDM(sw, &dm));
9199566063dSJacob Faibussowitsch     PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd));
920ea7032a0SMatthew G. Knepley     PetscCall(DMGetApplicationContext(sw, &ctx));
92135b38c60SMatthew G. Knepley 
92235b38c60SMatthew G. Knepley     /* Set particle position randomly in cell, set weights to 1 */
9239566063dSJacob Faibussowitsch     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)dm), &rnd));
9249566063dSJacob Faibussowitsch     PetscCall(PetscRandomSetInterval(rnd, -1.0, 1.0));
9259566063dSJacob Faibussowitsch     PetscCall(PetscRandomSetFromOptions(rnd));
9269566063dSJacob Faibussowitsch     PetscCall(DMSwarmSortGetAccess(sw));
92735b38c60SMatthew G. Knepley     for (d = 0; d < dim; ++d) xi0[d] = -1.0;
92835b38c60SMatthew G. Knepley     for (c = cStart; c < cEnd; ++c) {
92935b38c60SMatthew G. Knepley       PetscReal v0[3], J[9], invJ[9], detJ;
93035b38c60SMatthew G. Knepley       PetscInt *pidx, Npc, q;
93135b38c60SMatthew G. Knepley 
9329566063dSJacob Faibussowitsch       PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx));
9339566063dSJacob Faibussowitsch       PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v0, J, invJ, &detJ));
93435b38c60SMatthew G. Knepley       for (q = 0; q < Npc; ++q) {
93535b38c60SMatthew G. Knepley         const PetscInt p = pidx[q];
93635b38c60SMatthew G. Knepley         PetscReal      xref[3];
93735b38c60SMatthew G. Knepley 
9389566063dSJacob Faibussowitsch         for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &xref[d]));
93935b38c60SMatthew G. Knepley         CoordinatesRefToReal(dim, dim, xi0, v0, J, xref, &x[p * dim]);
94035b38c60SMatthew G. Knepley 
941ea7032a0SMatthew G. Knepley         weight[p]  = 1.0 / Np;
9426c5a79ebSMatthew G. Knepley         species[p] = p % Ns;
94335b38c60SMatthew G. Knepley       }
9449566063dSJacob Faibussowitsch       PetscCall(PetscFree(pidx));
94535b38c60SMatthew G. Knepley     }
9469566063dSJacob Faibussowitsch     PetscCall(PetscRandomDestroy(&rnd));
9479566063dSJacob Faibussowitsch     PetscCall(DMSwarmSortRestoreAccess(sw));
9486c5a79ebSMatthew G. Knepley   }
9499566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x));
9506c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight));
9519566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
9526c5a79ebSMatthew G. Knepley 
9539566063dSJacob Faibussowitsch   PetscCall(DMSwarmMigrate(sw, removePoints));
9549566063dSJacob Faibussowitsch   PetscCall(DMLocalizeCoordinates(sw));
9553ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
95635b38c60SMatthew G. Knepley }
95735b38c60SMatthew G. Knepley 
95835b38c60SMatthew G. Knepley /*@C
95935b38c60SMatthew G. Knepley   DMSwarmInitializeVelocities - Set the initial velocities of particles using a distribution.
96035b38c60SMatthew G. Knepley 
961*20f4b53cSBarry Smith   Collective
96235b38c60SMatthew G. Knepley 
96335b38c60SMatthew G. Knepley   Input Parameters:
964*20f4b53cSBarry Smith + sw      - The `DMSWARM` object
96535b38c60SMatthew G. Knepley . sampler - A function which uniformly samples the velocity PDF
96635b38c60SMatthew G. Knepley - v0      - The velocity scale for nondimensionalization for each species
96735b38c60SMatthew G. Knepley 
96835b38c60SMatthew G. Knepley   Level: advanced
96935b38c60SMatthew G. Knepley 
970*20f4b53cSBarry Smith   Note:
971*20f4b53cSBarry 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.
972*20f4b53cSBarry Smith 
973*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocitiesFromOptions()`
97435b38c60SMatthew G. Knepley @*/
975d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmInitializeVelocities(DM sw, PetscProbFunc sampler, const PetscReal v0[])
976d71ae5a4SJacob Faibussowitsch {
9776c5a79ebSMatthew G. Knepley   PetscSimplePointFunc velFunc;
97835b38c60SMatthew G. Knepley   PetscReal           *v;
97935b38c60SMatthew G. Knepley   PetscInt            *species;
9806c5a79ebSMatthew G. Knepley   void                *ctx;
98135b38c60SMatthew G. Knepley   PetscInt             dim, Np, p;
98235b38c60SMatthew G. Knepley 
98335b38c60SMatthew G. Knepley   PetscFunctionBegin;
9846c5a79ebSMatthew G. Knepley   PetscCall(DMSwarmGetVelocityFunction(sw, &velFunc));
98535b38c60SMatthew G. Knepley 
9869566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(sw, &dim));
9879566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetLocalSize(sw, &Np));
9889566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v));
9899566063dSJacob Faibussowitsch   PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species));
9906c5a79ebSMatthew G. Knepley   if (v0[0] == 0.) {
9916c5a79ebSMatthew G. Knepley     PetscCall(PetscArrayzero(v, Np * dim));
9926c5a79ebSMatthew G. Knepley   } else if (velFunc) {
9936c5a79ebSMatthew G. Knepley     PetscCall(DMGetApplicationContext(sw, &ctx));
9946c5a79ebSMatthew G. Knepley     for (p = 0; p < Np; ++p) {
9956c5a79ebSMatthew G. Knepley       PetscInt    s = species[p], d;
9966c5a79ebSMatthew G. Knepley       PetscScalar vel[3];
9976c5a79ebSMatthew G. Knepley 
9986c5a79ebSMatthew G. Knepley       PetscCall((*velFunc)(dim, 0., NULL, p, vel, ctx));
9996c5a79ebSMatthew G. Knepley       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * PetscRealPart(vel[d]);
10006c5a79ebSMatthew G. Knepley     }
10016c5a79ebSMatthew G. Knepley   } else {
10026c5a79ebSMatthew G. Knepley     PetscRandom rnd;
10036c5a79ebSMatthew G. Knepley 
10046c5a79ebSMatthew G. Knepley     PetscCall(PetscRandomCreate(PetscObjectComm((PetscObject)sw), &rnd));
10056c5a79ebSMatthew G. Knepley     PetscCall(PetscRandomSetInterval(rnd, 0, 1.));
10066c5a79ebSMatthew G. Knepley     PetscCall(PetscRandomSetFromOptions(rnd));
10076c5a79ebSMatthew G. Knepley 
100835b38c60SMatthew G. Knepley     for (p = 0; p < Np; ++p) {
100935b38c60SMatthew G. Knepley       PetscInt  s = species[p], d;
101035b38c60SMatthew G. Knepley       PetscReal a[3], vel[3];
101135b38c60SMatthew G. Knepley 
10129566063dSJacob Faibussowitsch       for (d = 0; d < dim; ++d) PetscCall(PetscRandomGetValueReal(rnd, &a[d]));
10139566063dSJacob Faibussowitsch       PetscCall(sampler(a, NULL, vel));
1014ad540459SPierre Jolivet       for (d = 0; d < dim; ++d) v[p * dim + d] = (v0[s] / v0[0]) * vel[d];
101535b38c60SMatthew G. Knepley     }
10166c5a79ebSMatthew G. Knepley     PetscCall(PetscRandomDestroy(&rnd));
10176c5a79ebSMatthew G. Knepley   }
10189566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v));
10199566063dSJacob Faibussowitsch   PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species));
10203ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
102135b38c60SMatthew G. Knepley }
102235b38c60SMatthew G. Knepley 
102335b38c60SMatthew G. Knepley /*@
102435b38c60SMatthew G. Knepley   DMSwarmInitializeVelocitiesFromOptions - Set the initial velocities of particles using a distribution determined from options.
102535b38c60SMatthew G. Knepley 
1026*20f4b53cSBarry Smith   Collective
102735b38c60SMatthew G. Knepley 
102835b38c60SMatthew G. Knepley   Input Parameters:
1029*20f4b53cSBarry Smith + sw      - The `DMSWARM` object
103035b38c60SMatthew G. Knepley - v0      - The velocity scale for nondimensionalization for each species
103135b38c60SMatthew G. Knepley 
103235b38c60SMatthew G. Knepley   Level: advanced
103335b38c60SMatthew G. Knepley 
1034*20f4b53cSBarry Smith .seealso: `DMSWARM`, `DMSwarmComputeLocalSize()`, `DMSwarmInitializeCoordinates()`, `DMSwarmInitializeVelocities()`
103535b38c60SMatthew G. Knepley @*/
1036d71ae5a4SJacob Faibussowitsch PetscErrorCode DMSwarmInitializeVelocitiesFromOptions(DM sw, const PetscReal v0[])
1037d71ae5a4SJacob Faibussowitsch {
103835b38c60SMatthew G. Knepley   PetscProbFunc sampler;
103935b38c60SMatthew G. Knepley   PetscInt      dim;
104035b38c60SMatthew G. Knepley   const char   *prefix;
10416c5a79ebSMatthew G. Knepley   char          funcname[PETSC_MAX_PATH_LEN];
10426c5a79ebSMatthew G. Knepley   PetscBool     flg;
104335b38c60SMatthew G. Knepley 
104435b38c60SMatthew G. Knepley   PetscFunctionBegin;
1045d0609cedSBarry Smith   PetscOptionsBegin(PetscObjectComm((PetscObject)sw), "", "DMSwarm Options", "DMSWARM");
10466c5a79ebSMatthew G. Knepley   PetscCall(PetscOptionsString("-dm_swarm_velocity_function", "Function to determine particle velocities", "DMSwarmSetVelocityFunction", funcname, funcname, sizeof(funcname), &flg));
1047d0609cedSBarry Smith   PetscOptionsEnd();
10486c5a79ebSMatthew G. Knepley   if (flg) {
10496c5a79ebSMatthew G. Knepley     PetscSimplePointFunc velFunc;
10506c5a79ebSMatthew G. Knepley 
10516c5a79ebSMatthew G. Knepley     PetscCall(PetscDLSym(NULL, funcname, (void **)&velFunc));
10526c5a79ebSMatthew G. Knepley     PetscCheck(velFunc, PetscObjectComm((PetscObject)sw), PETSC_ERR_ARG_WRONG, "Could not locate function %s", funcname);
10536c5a79ebSMatthew G. Knepley     PetscCall(DMSwarmSetVelocityFunction(sw, velFunc));
10546c5a79ebSMatthew G. Knepley   }
10559566063dSJacob Faibussowitsch   PetscCall(DMGetDimension(sw, &dim));
10569566063dSJacob Faibussowitsch   PetscCall(PetscObjectGetOptionsPrefix((PetscObject)sw, &prefix));
10579566063dSJacob Faibussowitsch   PetscCall(PetscProbCreateFromOptions(dim, prefix, "-dm_swarm_velocity_density", NULL, NULL, &sampler));
10589566063dSJacob Faibussowitsch   PetscCall(DMSwarmInitializeVelocities(sw, sampler, v0));
10593ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
106035b38c60SMatthew G. Knepley }
1061