Actual source code: ex7.c

  1: static char help[] = "Example program demonstrating projection between particle and finite element spaces using OpenMP in 2D cylindrical coordinates\n";

  3: #include "petscdmplex.h"
  4: #include "petscds.h"
  5: #include "petscdmswarm.h"
  6: #include "petscksp.h"
  7: #include <petsc/private/petscimpl.h>
  8: #if defined(PETSC_HAVE_OPENMP) && defined(PETSC_HAVE_THREADSAFETY)
  9:   #include <omp.h>
 10: #endif

 12: typedef struct {
 13:   Mat MpTrans;
 14:   Mat Mp;
 15:   Vec ff;
 16:   Vec uu;
 17: } MatShellCtx;

 19: PetscErrorCode MatMultMtM_SeqAIJ(Mat MtM, Vec xx, Vec yy)
 20: {
 21:   MatShellCtx *matshellctx;

 24:   MatShellGetContext(MtM, &matshellctx);
 26:   MatMult(matshellctx->Mp, xx, matshellctx->ff);
 27:   MatMult(matshellctx->MpTrans, matshellctx->ff, yy);
 28:   return 0;
 29: }

 31: PetscErrorCode MatMultAddMtM_SeqAIJ(Mat MtM, Vec xx, Vec yy, Vec zz)
 32: {
 33:   MatShellCtx *matshellctx;

 36:   MatShellGetContext(MtM, &matshellctx);
 38:   MatMult(matshellctx->Mp, xx, matshellctx->ff);
 39:   MatMultAdd(matshellctx->MpTrans, matshellctx->ff, yy, zz);
 40:   return 0;
 41: }

 43: PetscErrorCode createSwarm(const DM dm, DM *sw)
 44: {
 45:   PetscInt Nc = 1, dim = 2;

 48:   DMCreate(PETSC_COMM_SELF, sw);
 49:   DMSetType(*sw, DMSWARM);
 50:   DMSetDimension(*sw, dim);
 51:   DMSwarmSetType(*sw, DMSWARM_PIC);
 52:   DMSwarmSetCellDM(*sw, dm);
 53:   DMSwarmRegisterPetscDatatypeField(*sw, "w_q", Nc, PETSC_SCALAR);
 54:   DMSwarmFinalizeFieldRegister(*sw);
 55:   DMSetFromOptions(*sw);
 56:   return 0;
 57: }

 59: PetscErrorCode gridToParticles(const DM dm, DM sw, PetscReal *moments, Vec rhs, Mat M_p)
 60: {
 61:   PetscBool     is_lsqr;
 62:   KSP           ksp;
 63:   Mat           PM_p = NULL, MtM, D;
 64:   Vec           ff;
 65:   PetscInt      Np, timestep = 0, bs, N, M, nzl;
 66:   PetscReal     time = 0.0;
 67:   PetscDataType dtype;
 68:   MatShellCtx  *matshellctx;

 71:   KSPCreate(PETSC_COMM_SELF, &ksp);
 72:   KSPSetOptionsPrefix(ksp, "ftop_");
 73:   KSPSetFromOptions(ksp);
 74:   PetscObjectTypeCompare((PetscObject)ksp, KSPLSQR, &is_lsqr);
 75:   if (!is_lsqr) {
 76:     MatGetLocalSize(M_p, &M, &N);
 77:     if (N > M) {
 78:       PC pc;
 79:       PetscInfo(ksp, " M (%" PetscInt_FMT ") < M (%" PetscInt_FMT ") -- skip revert to lsqr\n", M, N);
 80:       is_lsqr = PETSC_TRUE;
 81:       KSPSetType(ksp, KSPLSQR);
 82:       KSPGetPC(ksp, &pc);
 83:       PCSetType(pc, PCNONE); // could put in better solver -ftop_pc_type bjacobi -ftop_sub_pc_type lu -ftop_sub_pc_factor_shift_type nonzero
 84:     } else {
 85:       PetscNew(&matshellctx);
 86:       MatCreateShell(PetscObjectComm((PetscObject)dm), N, N, PETSC_DECIDE, PETSC_DECIDE, matshellctx, &MtM);
 87:       MatTranspose(M_p, MAT_INITIAL_MATRIX, &matshellctx->MpTrans);
 88:       matshellctx->Mp = M_p;
 89:       MatShellSetOperation(MtM, MATOP_MULT, (void (*)(void))MatMultMtM_SeqAIJ);
 90:       MatShellSetOperation(MtM, MATOP_MULT_ADD, (void (*)(void))MatMultAddMtM_SeqAIJ);
 91:       MatCreateVecs(M_p, &matshellctx->uu, &matshellctx->ff);
 92:       MatCreateSeqAIJ(PETSC_COMM_SELF, N, N, 1, NULL, &D);
 93:       for (int i = 0; i < N; i++) {
 94:         const PetscScalar *vals;
 95:         const PetscInt    *cols;
 96:         PetscScalar        dot = 0;
 97:         MatGetRow(matshellctx->MpTrans, i, &nzl, &cols, &vals);
 98:         for (int ii = 0; ii < nzl; ii++) dot += PetscSqr(vals[ii]);
100:         MatSetValue(D, i, i, dot, INSERT_VALUES);
101:       }
102:       MatAssemblyBegin(D, MAT_FINAL_ASSEMBLY);
103:       MatAssemblyEnd(D, MAT_FINAL_ASSEMBLY);
104:       PetscInfo(M_p, "createMtMKSP Have %" PetscInt_FMT " eqs, nzl = %" PetscInt_FMT "\n", N, nzl);
105:       KSPSetOperators(ksp, MtM, D);
106:       MatViewFromOptions(D, NULL, "-ftop2_D_mat_view");
107:       MatViewFromOptions(M_p, NULL, "-ftop2_Mp_mat_view");
108:       MatViewFromOptions(matshellctx->MpTrans, NULL, "-ftop2_MpTranspose_mat_view");
109:     }
110:   }
111:   if (is_lsqr) {
112:     PC        pc;
113:     PetscBool is_bjac;
114:     KSPGetPC(ksp, &pc);
115:     PetscObjectTypeCompare((PetscObject)pc, PCBJACOBI, &is_bjac);
116:     if (is_bjac) {
117:       DMSwarmCreateMassMatrixSquare(sw, dm, &PM_p);
118:       KSPSetOperators(ksp, M_p, PM_p);
119:     } else {
120:       KSPSetOperators(ksp, M_p, M_p);
121:     }
122:   }
123:   DMSwarmCreateGlobalVectorFromField(sw, "w_q", &ff); // this grabs access !!!!!
124:   if (!is_lsqr) {
125:     KSPSolve(ksp, rhs, matshellctx->uu);
126:     MatMult(M_p, matshellctx->uu, ff);
127:     MatDestroy(&matshellctx->MpTrans);
128:     VecDestroy(&matshellctx->ff);
129:     VecDestroy(&matshellctx->uu);
130:     MatDestroy(&D);
131:     MatDestroy(&MtM);
132:     PetscFree(matshellctx);
133:   } else {
134:     KSPSolveTranspose(ksp, rhs, ff);
135:   }
136:   KSPDestroy(&ksp);
137:   /* Visualize particle field */
138:   DMSetOutputSequenceNumber(sw, timestep, time);
139:   VecViewFromOptions(ff, NULL, "-weights_view");
140:   DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &ff);

142:   /* compute energy */
143:   if (moments) {
144:     PetscReal *wq, *coords;
145:     DMSwarmGetLocalSize(sw, &Np);
146:     DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wq);
147:     DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords);
148:     moments[0] = moments[1] = moments[2] = 0;
149:     for (int p = 0; p < Np; p++) {
150:       moments[0] += wq[p];
151:       moments[1] += wq[p] * coords[p * 2 + 0]; // x-momentum
152:       moments[2] += wq[p] * (PetscSqr(coords[p * 2 + 0]) + PetscSqr(coords[p * 2 + 1]));
153:     }
154:     DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords);
155:     DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wq);
156:   }
157:   MatDestroy(&PM_p);
158:   return 0;
159: }

161: PetscErrorCode particlesToGrid(const DM dm, DM sw, const PetscInt Np, const PetscInt a_tid, const PetscInt dim, const PetscInt target, const PetscReal xx[], const PetscReal yy[], const PetscReal a_wp[], Vec rho, Mat *Mp_out)
162: {
163:   PetscBool     removePoints = PETSC_TRUE;
164:   PetscReal    *wq, *coords;
165:   PetscDataType dtype;
166:   Mat           M_p;
167:   Vec           ff;
168:   PetscInt      bs, p, zero = 0;

171:   DMSwarmSetLocalSizes(sw, Np, zero);
172:   DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wq);
173:   DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords);
174:   for (p = 0; p < Np; p++) {
175:     coords[p * 2 + 0] = xx[p];
176:     coords[p * 2 + 1] = yy[p];
177:     wq[p]             = a_wp[p];
178:   }
179:   DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords);
180:   DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wq);
181:   DMSwarmMigrate(sw, removePoints);
182:   PetscObjectSetName((PetscObject)sw, "Particle Grid");
183:   if (a_tid == target) DMViewFromOptions(sw, NULL, "-swarm_view");

185:   /* Project particles to field */
186:   /* This gives M f = \int_\Omega \phi f, which looks like a rhs for a PDE */
187:   DMCreateMassMatrix(sw, dm, &M_p);
188:   PetscObjectSetName((PetscObject)rho, "rho");

190:   DMSwarmCreateGlobalVectorFromField(sw, "w_q", &ff); // this grabs access !!!!!
191:   PetscObjectSetName((PetscObject)ff, "weights");
192:   MatMultTranspose(M_p, ff, rho);
193:   DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &ff);

195:   /* Visualize mesh field */
196:   if (a_tid == target) VecViewFromOptions(rho, NULL, "-rho_view");
197:   // output
198:   *Mp_out = M_p;

200:   return 0;
201: }
202: static PetscErrorCode maxwellian(PetscInt dim, const PetscReal x[], PetscReal kt_m, PetscReal n, PetscScalar *u)
203: {
204:   PetscInt  i;
205:   PetscReal v2 = 0, theta = 2 * kt_m; /* theta = 2kT/mc^2 */

207:   /* compute the exponents, v^2 */
208:   for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
209:   /* evaluate the Maxwellian */
210:   u[0] = n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta)) * 2. * PETSC_PI * x[1]; // radial term for 2D axi-sym.
211:   return 0;
212: }

214: #define MAX_NUM_THRDS 12
215: PetscErrorCode go()
216: {
217:   DM        dm_t[MAX_NUM_THRDS], sw_t[MAX_NUM_THRDS];
218:   PetscFE   fe;
219:   PetscInt  dim = 2, Nc = 1, i, faces[3];
220:   PetscInt  Np[2] = {10, 10}, Np2[2], field = 0, target = 0, Np_t[MAX_NUM_THRDS];
221:   PetscReal moments_0[3], moments_1[3], vol = 1;
222:   PetscReal lo[3] = {-5, 0, -5}, hi[3] = {5, 5, 5}, h[3], hp[3], *xx_t[MAX_NUM_THRDS], *yy_t[MAX_NUM_THRDS], *wp_t[MAX_NUM_THRDS];
223:   Vec       rho_t[MAX_NUM_THRDS], rhs_t[MAX_NUM_THRDS];
224:   Mat       M_p_t[MAX_NUM_THRDS];
225: #if defined PETSC_USE_LOG
226:   PetscLogStage stage;
227:   PetscLogEvent swarm_create_ev, solve_ev, solve_loop_ev;
228: #endif
229: #if defined(PETSC_HAVE_OPENMP) && defined(PETSC_HAVE_THREADSAFETY)
230:   PetscInt numthreads = PetscNumOMPThreads;
231: #else
232:   PetscInt numthreads = 1;
233: #endif

236: #if defined(PETSC_HAVE_OPENMP) && defined(PETSC_HAVE_THREADSAFETY)
239: #endif
240:   if (target >= numthreads) target = numthreads - 1;
241:   PetscLogEventRegister("Create Swarm", DM_CLASSID, &swarm_create_ev);
242:   PetscLogEventRegister("Single solve", DM_CLASSID, &solve_ev);
243:   PetscLogEventRegister("Solve loop", DM_CLASSID, &solve_loop_ev);
244:   PetscLogStageRegister("Solve", &stage);
245:   i = dim;
246:   PetscOptionsGetIntArray(NULL, NULL, "-dm_plex_box_faces", faces, &i, NULL);
247:   i = dim;
248:   PetscOptionsGetIntArray(NULL, NULL, "-np", Np, &i, NULL);
249:   /* Create thread meshes */
250:   for (int tid = 0; tid < numthreads; tid++) {
251:     // setup mesh dm_t, could use PETSc's Landau create velocity space mesh here to get dm_t[tid]
252:     DMCreate(PETSC_COMM_SELF, &dm_t[tid]);
253:     DMSetType(dm_t[tid], DMPLEX);
254:     DMSetFromOptions(dm_t[tid]);
255:     PetscFECreateDefault(PETSC_COMM_SELF, dim, Nc, PETSC_FALSE, "", PETSC_DECIDE, &fe);
256:     PetscFESetFromOptions(fe);
257:     PetscObjectSetName((PetscObject)fe, "fe");
258:     DMSetField(dm_t[tid], field, NULL, (PetscObject)fe);
259:     DMCreateDS(dm_t[tid]);
260:     PetscFEDestroy(&fe);
261:     // helper vectors
262:     DMCreateGlobalVector(dm_t[tid], &rho_t[tid]);
263:     DMCreateGlobalVector(dm_t[tid], &rhs_t[tid]);
264:     // this mimics application code
265:     DMGetBoundingBox(dm_t[tid], lo, hi);
266:     if (tid == target) {
267:       DMViewFromOptions(dm_t[tid], NULL, "-dm_view");
268:       for (i = 0, vol = 1; i < dim; i++) {
269:         h[i]  = (hi[i] - lo[i]) / faces[i];
270:         hp[i] = (hi[i] - lo[i]) / Np[i];
271:         vol *= (hi[i] - lo[i]);
272:         PetscInfo(dm_t[tid], " lo = %g hi = %g n = %" PetscInt_FMT " h = %g hp = %g\n", (double)lo[i], (double)hi[i], faces[i], (double)h[i], (double)hp[i]);
273:       }
274:     }
275:   }
276:   // prepare particle data for problems. This mimics application code
277:   PetscLogEventBegin(swarm_create_ev, 0, 0, 0, 0);
278:   Np2[0] = Np[0];
279:   Np2[1] = Np[1];
280:   for (int tid = 0; tid < numthreads; tid++) { // change size of particle list a little
281:     Np_t[tid] = Np2[0] * Np2[1];
282:     PetscMalloc3(Np_t[tid], &xx_t[tid], Np_t[tid], &yy_t[tid], Np_t[tid], &wp_t[tid]);
283:     if (tid == target) moments_0[0] = moments_0[1] = moments_0[2] = 0;
284:     for (int pi = 0, pp = 0; pi < Np2[0]; pi++) {
285:       for (int pj = 0; pj < Np2[1]; pj++, pp++) {
286:         xx_t[tid][pp] = lo[0] + hp[0] / 2. + pi * hp[0];
287:         yy_t[tid][pp] = lo[1] + hp[1] / 2. + pj * hp[1];
288:         {
289:           PetscReal x[] = {xx_t[tid][pp], yy_t[tid][pp]};
290:           maxwellian(2, x, 1.0, vol / (PetscReal)Np_t[tid], &wp_t[tid][pp]);
291:         }
292:         if (tid == target) { //energy_0 += wp_t[tid][pp]*(PetscSqr(xx_t[tid][pp])+PetscSqr(yy_t[tid][pp]));
293:           moments_0[0] += wp_t[tid][pp];
294:           moments_0[1] += wp_t[tid][pp] * xx_t[tid][pp]; // x-momentum
295:           moments_0[2] += wp_t[tid][pp] * (PetscSqr(xx_t[tid][pp]) + PetscSqr(yy_t[tid][pp]));
296:         }
297:       }
298:     }
299:     Np2[0]++;
300:     Np2[1]++;
301:   }
302:   PetscLogEventEnd(swarm_create_ev, 0, 0, 0, 0);
303:   PetscLogEventBegin(solve_ev, 0, 0, 0, 0);
304:   /* Create particle swarm */
305:   PetscPragmaOMP(parallel for)
306:   for (int tid=0; tid<numthreads; tid++)
307:   {
308:     PETSC_COMM_SELF, createSwarm(dm_t[tid], &sw_t[tid]);
309:   }
310:   PetscPragmaOMP(parallel for)
311:   for (int tid=0; tid<numthreads; tid++)
312:   {
313:     PETSC_COMM_SELF, particlesToGrid(dm_t[tid], sw_t[tid], Np_t[tid], tid, dim, target, xx_t[tid], yy_t[tid], wp_t[tid], rho_t[tid], &M_p_t[tid]);
314:   }
315:   /* Project field to particles */
316:   /*   This gives f_p = M_p^+ M f */
317:   PetscPragmaOMP(parallel for)
318:   for (int tid=0; tid<numthreads; tid++)
319:   {
320:     PETSC_COMM_SELF, VecCopy(rho_t[tid], rhs_t[tid]); /* Identity: M^1 M rho */
321:   }
322:   PetscPragmaOMP(parallel for)
323:   for (int tid=0; tid<numthreads; tid++)
324:   {
325:     PETSC_COMM_SELF, gridToParticles(dm_t[tid], sw_t[tid], (tid == target) ? moments_1 : NULL, rhs_t[tid], M_p_t[tid]);
326:   }
327:   /* Cleanup */
328:   for (int tid = 0; tid < numthreads; tid++) {
329:     MatDestroy(&M_p_t[tid]);
330:     DMDestroy(&sw_t[tid]);
331:   }
332:   PetscLogEventEnd(solve_ev, 0, 0, 0, 0);
333:   //
334:   PetscPrintf(PETSC_COMM_SELF, "Total number density: %20.12e (%20.12e); x-momentum = %g (%g); energy = %g error = %e, %" PetscInt_FMT " particles. Use %" PetscInt_FMT " threads\n", (double)moments_1[0], (double)moments_0[0], (double)moments_1[1], (double)moments_0[1], (double)moments_1[2], (double)((moments_1[2] - moments_0[2]) / moments_0[2]), Np[0] * Np[1], numthreads);
335:   /* Cleanup */
336:   for (int tid = 0; tid < numthreads; tid++) {
337:     VecDestroy(&rho_t[tid]);
338:     VecDestroy(&rhs_t[tid]);
339:     DMDestroy(&dm_t[tid]);
340:     PetscFree3(xx_t[tid], yy_t[tid], wp_t[tid]);
341:   }
342:   return 0;
343: }

345: int main(int argc, char **argv)
346: {
348:   PetscInitialize(&argc, &argv, NULL, help);
349:   go();
350:   PetscFinalize();
351:   return 0;
352: }

354: /*TEST

356:   build:
357:     requires: !complex

359:   test:
360:     suffix: 0
361:     requires: double triangle
362:     args: -dm_plex_simplex 0 -dm_plex_box_faces 8,4 -np 10 -dm_plex_box_lower -2.0,0.0 -dm_plex_box_upper 2.0,2.0 -petscspace_degree 2 -ftop_ksp_type lsqr -ftop_pc_type none -dm_view -ftop_ksp_converged_reason -ftop_ksp_rtol 1.e-14
363:     filter: grep -v DM_ | grep -v atomic

365:   test:
366:     suffix: 1
367:     requires: double triangle
368:     args: -dm_plex_simplex 0 -dm_plex_box_faces 8,4 -np 10 -dm_plex_box_lower -2.0,0.0 -dm_plex_box_upper 2.0,2.0 -petscspace_degree 2 -dm_plex_hash_location -ftop_ksp_type lsqr -ftop_pc_type bjacobi -ftop_sub_pc_type lu -ftop_sub_pc_factor_shift_type nonzero -dm_view -ftop_ksp_converged_reason -ftop_ksp_rtol 1.e-14
369:     filter: grep -v DM_ | grep -v atomic

371:   test:
372:     suffix: 2
373:     requires: double triangle
374:     args: -dm_plex_simplex 0 -dm_plex_box_faces 8,4 -np 10 -dm_plex_box_lower -2.0,0.0 -dm_plex_box_upper 2.0,2.0 -petscspace_degree 2 -dm_plex_hash_location -ftop_ksp_type cg -ftop_pc_type jacobi -dm_view -ftop_ksp_converged_reason -ftop_ksp_rtol 1.e-14
375:     filter: grep -v DM_ | grep -v atomic

377: TEST*/