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*/