Actual source code: ex30.c
1: static char help[] = "Grid based Landau collision operator with PIC interface with OpenMP setup. (one species per grid)\n";
3: /*
4: Support 2.5V with axisymmetric coordinates
5: - r,z coordinates
6: - Domain and species data input by Landau operator
7: - "radius" for each grid, normalized with electron thermal velocity
8: - Domain: (0,radius) x (-radius,radius), thus first coordinate x[0] is perpendicular velocity and 2pi*x[0] term is added for axisymmetric
9: Supports full 3V
11: */
13: #include "petscdmplex.h"
14: #include "petscds.h"
15: #include "petscdmswarm.h"
16: #include "petscksp.h"
17: #include <petsc/private/petscimpl.h>
18: #if defined(PETSC_HAVE_OPENMP) && defined(PETSC_HAVE_THREADSAFETY)
19: #include <omp.h>
20: #endif
21: #include <petsclandau.h>
22: #include <petscdmcomposite.h>
24: typedef struct {
25: Mat MpTrans;
26: Mat Mp;
27: Vec ff;
28: Vec uu;
29: } MatShellCtx;
31: PetscErrorCode MatMultMtM_SeqAIJ(Mat MtM, Vec xx, Vec yy)
32: {
33: MatShellCtx *matshellctx;
35: PetscFunctionBeginUser;
36: PetscCall(MatShellGetContext(MtM, &matshellctx));
37: PetscCheck(matshellctx, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "No context");
38: PetscCall(MatMult(matshellctx->Mp, xx, matshellctx->ff));
39: PetscCall(MatMult(matshellctx->MpTrans, matshellctx->ff, yy));
40: PetscFunctionReturn(PETSC_SUCCESS);
41: }
43: PetscErrorCode MatMultAddMtM_SeqAIJ(Mat MtM, Vec xx, Vec yy, Vec zz)
44: {
45: MatShellCtx *matshellctx;
47: PetscFunctionBeginUser;
48: PetscCall(MatShellGetContext(MtM, &matshellctx));
49: PetscCheck(matshellctx, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "No context");
50: PetscCall(MatMult(matshellctx->Mp, xx, matshellctx->ff));
51: PetscCall(MatMultAdd(matshellctx->MpTrans, matshellctx->ff, yy, zz));
52: PetscFunctionReturn(PETSC_SUCCESS);
53: }
55: PetscErrorCode createSwarm(const DM dm, PetscInt dim, DM *sw)
56: {
57: PetscInt Nc = 1;
59: PetscFunctionBeginUser;
60: PetscCall(DMCreate(PETSC_COMM_SELF, sw));
61: PetscCall(DMSetType(*sw, DMSWARM));
62: PetscCall(DMSetDimension(*sw, dim));
63: PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
64: PetscCall(DMSwarmSetCellDM(*sw, dm));
65: PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", Nc, PETSC_SCALAR));
66: PetscCall(DMSwarmFinalizeFieldRegister(*sw));
67: PetscCall(DMSetFromOptions(*sw));
68: PetscFunctionReturn(PETSC_SUCCESS);
69: }
71: PetscErrorCode gridToParticles(const DM dm, DM sw, Vec rhs, Vec work, Mat M_p, Mat Mass)
72: {
73: PetscBool is_lsqr;
74: KSP ksp;
75: Mat PM_p = NULL, MtM, D;
76: Vec ff;
77: PetscInt N, M, nzl;
78: MatShellCtx *matshellctx;
80: PetscFunctionBeginUser;
81: PetscCall(MatMult(Mass, rhs, work));
82: PetscCall(VecCopy(work, rhs));
83: // pseudo-inverse
84: PetscCall(KSPCreate(PETSC_COMM_SELF, &ksp));
85: PetscCall(KSPSetOptionsPrefix(ksp, "ftop_"));
86: PetscCall(KSPSetFromOptions(ksp));
87: PetscCall(PetscObjectTypeCompare((PetscObject)ksp, KSPLSQR, &is_lsqr));
88: if (!is_lsqr) {
89: PetscCall(MatGetLocalSize(M_p, &M, &N));
90: if (N > M) {
91: PC pc;
92: PetscCall(PetscInfo(ksp, " M (%" PetscInt_FMT ") < M (%" PetscInt_FMT ") -- skip revert to lsqr\n", M, N));
93: is_lsqr = PETSC_TRUE;
94: PetscCall(KSPSetType(ksp, KSPLSQR));
95: PetscCall(KSPGetPC(ksp, &pc));
96: PetscCall(PCSetType(pc, PCNONE)); // could put in better solver -ftop_pc_type bjacobi -ftop_sub_pc_type lu -ftop_sub_pc_factor_shift_type nonzero
97: } else {
98: PetscCall(PetscNew(&matshellctx));
99: PetscCall(MatCreateShell(PetscObjectComm((PetscObject)dm), N, N, PETSC_DECIDE, PETSC_DECIDE, matshellctx, &MtM));
100: PetscCall(MatTranspose(M_p, MAT_INITIAL_MATRIX, &matshellctx->MpTrans));
101: matshellctx->Mp = M_p;
102: PetscCall(MatShellSetOperation(MtM, MATOP_MULT, (void (*)(void))MatMultMtM_SeqAIJ));
103: PetscCall(MatShellSetOperation(MtM, MATOP_MULT_ADD, (void (*)(void))MatMultAddMtM_SeqAIJ));
104: PetscCall(MatCreateVecs(M_p, &matshellctx->uu, &matshellctx->ff));
105: PetscCall(MatCreateSeqAIJ(PETSC_COMM_SELF, N, N, 1, NULL, &D));
106: PetscCall(MatViewFromOptions(matshellctx->MpTrans, NULL, "-ftop2_Mp_mat_view"));
107: for (int i = 0; i < N; i++) {
108: const PetscScalar *vals;
109: const PetscInt *cols;
110: PetscScalar dot = 0;
111: PetscCall(MatGetRow(matshellctx->MpTrans, i, &nzl, &cols, &vals));
112: for (int ii = 0; ii < nzl; ii++) dot += PetscSqr(vals[ii]);
113: PetscCheck(dot != 0.0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Row %d is empty", i);
114: PetscCall(MatSetValue(D, i, i, dot, INSERT_VALUES));
115: }
116: PetscCall(MatAssemblyBegin(D, MAT_FINAL_ASSEMBLY));
117: PetscCall(MatAssemblyEnd(D, MAT_FINAL_ASSEMBLY));
118: PetscCall(PetscInfo(M_p, "createMtMKSP Have %" PetscInt_FMT " eqs, nzl = %" PetscInt_FMT "\n", N, nzl));
119: PetscCall(KSPSetOperators(ksp, MtM, D));
120: PetscCall(MatViewFromOptions(D, NULL, "-ftop2_D_mat_view"));
121: PetscCall(MatViewFromOptions(M_p, NULL, "-ftop2_Mp_mat_view"));
122: PetscCall(MatViewFromOptions(matshellctx->MpTrans, NULL, "-ftop2_MpTranspose_mat_view"));
123: PetscCall(MatViewFromOptions(MtM, NULL, "-ftop2_MtM_mat_view"));
124: }
125: }
126: if (is_lsqr) {
127: PC pc;
128: PetscBool is_bjac;
129: PetscCall(KSPGetPC(ksp, &pc));
130: PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCBJACOBI, &is_bjac));
131: if (is_bjac) {
132: PetscCall(DMSwarmCreateMassMatrixSquare(sw, dm, &PM_p));
133: PetscCall(KSPSetOperators(ksp, M_p, PM_p));
134: } else {
135: PetscCall(KSPSetOperators(ksp, M_p, M_p));
136: }
137: }
138: PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &ff)); // this grabs access
139: if (!is_lsqr) {
140: PetscCall(KSPSolve(ksp, rhs, matshellctx->uu));
141: PetscCall(MatMult(M_p, matshellctx->uu, ff));
142: PetscCall(MatDestroy(&matshellctx->MpTrans));
143: PetscCall(VecDestroy(&matshellctx->ff));
144: PetscCall(VecDestroy(&matshellctx->uu));
145: PetscCall(MatDestroy(&D));
146: PetscCall(MatDestroy(&MtM));
147: PetscCall(PetscFree(matshellctx));
148: } else {
149: PetscCall(KSPSolveTranspose(ksp, rhs, ff));
150: }
151: PetscCall(KSPDestroy(&ksp));
152: /* Visualize particle field */
153: PetscCall(VecViewFromOptions(ff, NULL, "-weights_view"));
154: PetscCall(MatDestroy(&PM_p));
155: PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &ff));
157: PetscFunctionReturn(PETSC_SUCCESS);
158: }
160: PetscErrorCode particlesToGrid(const DM dm, DM sw, const PetscInt Np, const PetscInt a_tid, const PetscInt dim, const PetscReal xx[], const PetscReal yy[], const PetscReal zz[], const PetscReal a_wp[], Vec rho, Mat *Mp_out)
161: {
162: PetscBool removePoints = PETSC_TRUE;
163: PetscReal *wq, *coords;
164: PetscDataType dtype;
165: Mat M_p;
166: Vec ff;
167: PetscInt bs, p, zero = 0;
169: PetscFunctionBeginUser;
170: PetscCall(DMSwarmSetLocalSizes(sw, Np, zero));
171: PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wq));
172: PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
173: for (p = 0; p < Np; p++) {
174: coords[p * dim + 0] = xx[p];
175: coords[p * dim + 1] = yy[p];
176: wq[p] = a_wp[p];
177: if (dim == 3) coords[p * dim + 2] = zz[p];
178: }
179: PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
180: PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wq));
181: PetscCall(DMSwarmMigrate(sw, removePoints));
182: PetscCall(PetscObjectSetName((PetscObject)sw, "Particle Grid"));
184: /* This gives M f = \int_\Omega \phi f, which looks like a rhs for a PDE */
185: PetscCall(DMCreateMassMatrix(sw, dm, &M_p));
187: PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
188: PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &ff));
189: PetscCall(PetscObjectSetName((PetscObject)ff, "weights"));
190: PetscCall(MatMultTranspose(M_p, ff, rho));
191: PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &ff));
193: // output
194: *Mp_out = M_p;
196: PetscFunctionReturn(PETSC_SUCCESS);
197: }
198: static void maxwellian(PetscInt dim, const PetscReal x[], PetscReal kt_m, PetscReal n, PetscScalar *u)
199: {
200: PetscInt i;
201: PetscReal v2 = 0, theta = 2.0 * kt_m; /* theta = 2kT/mc^2 */
203: /* compute the exponents, v^2 */
204: for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
205: /* evaluate the Maxwellian */
206: u[0] = n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
207: }
209: #define MAX_NUM_THRDS 12
210: PetscErrorCode go(TS ts, Vec X, const PetscInt NUserV, const PetscInt a_Np, const PetscInt dim, const PetscInt b_target, const PetscInt g_target)
211: {
212: DM pack, *globSwarmArray, grid_dm[LANDAU_MAX_GRIDS];
213: Mat *globMpArray, g_Mass[LANDAU_MAX_GRIDS];
214: KSP t_ksp[LANDAU_MAX_GRIDS][MAX_NUM_THRDS];
215: Vec t_fhat[LANDAU_MAX_GRIDS][MAX_NUM_THRDS];
216: PetscInt nDMs, glb_b_id, nTargetP = 0;
217: PetscErrorCode ierr = 0; // used for inside thread loops
218: #if defined(PETSC_HAVE_OPENMP) && defined(PETSC_HAVE_THREADSAFETY)
219: PetscInt numthreads = PetscNumOMPThreads;
220: #else
221: PetscInt numthreads = 1;
222: #endif
223: LandauCtx *ctx;
224: Vec *globXArray;
225: PetscReal moments_0[3], moments_1[3], dt_init;
227: PetscFunctionBeginUser;
228: PetscCheck(numthreads <= MAX_NUM_THRDS, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Too many threads %" PetscInt_FMT " > %d", numthreads, MAX_NUM_THRDS);
229: PetscCheck(numthreads > 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number threads %" PetscInt_FMT " > %d", numthreads, MAX_NUM_THRDS);
230: PetscCall(TSGetDM(ts, &pack));
231: PetscCall(DMGetApplicationContext(pack, &ctx));
232: PetscCheck(ctx->batch_sz % numthreads == 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "batch size (-dm_landau_batch_size) %" PetscInt_FMT " mod #threads %" PetscInt_FMT " must equal zero", ctx->batch_sz, numthreads);
233: PetscCall(DMCompositeGetNumberDM(pack, &nDMs));
234: PetscCall(PetscInfo(pack, "Have %" PetscInt_FMT " total grids, with %" PetscInt_FMT " Landau local batched and %" PetscInt_FMT " global items (vertices)\n", ctx->num_grids, ctx->batch_sz, NUserV));
235: PetscCall(PetscMalloc(sizeof(*globXArray) * nDMs, &globXArray));
236: PetscCall(PetscMalloc(sizeof(*globMpArray) * nDMs, &globMpArray));
237: PetscCall(PetscMalloc(sizeof(*globSwarmArray) * nDMs, &globSwarmArray));
238: PetscCall(DMViewFromOptions(ctx->plex[g_target], NULL, "-ex30_dm_view"));
239: // create mass matrices
240: PetscCall(VecZeroEntries(X));
241: PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, globXArray)); // just to duplicate
242: for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
243: Vec subX = globXArray[LAND_PACK_IDX(0, grid)];
244: DM dm = ctx->plex[grid];
245: PetscSection s;
246: grid_dm[grid] = dm;
247: PetscCall(DMCreateMassMatrix(dm, dm, &g_Mass[grid]));
248: //
249: PetscCall(DMGetLocalSection(dm, &s));
250: PetscCall(DMPlexCreateClosureIndex(dm, s));
251: for (int tid = 0; tid < numthreads; tid++) {
252: PetscCall(VecDuplicate(subX, &t_fhat[grid][tid]));
253: PetscCall(KSPCreate(PETSC_COMM_SELF, &t_ksp[grid][tid]));
254: PetscCall(KSPSetOptionsPrefix(t_ksp[grid][tid], "ptof_"));
255: PetscCall(KSPSetOperators(t_ksp[grid][tid], g_Mass[grid], g_Mass[grid]));
256: PetscCall(KSPSetFromOptions(t_ksp[grid][tid]));
257: }
258: }
259: PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, globXArray));
260: // create particle raw data. could use OMP with a thread safe malloc, but this is just the fake user
261: for (int i = 0; i < 3; i++) moments_0[i] = moments_1[i] = 0;
262: PetscCall(TSGetTimeStep(ts, &dt_init)); // we could have an adaptive time stepper
263: for (PetscInt global_batch_id = 0; global_batch_id < NUserV; global_batch_id += ctx->batch_sz) {
264: PetscCall(TSSetTime(ts, 0));
265: PetscCall(TSSetStepNumber(ts, 0));
266: PetscCall(TSSetTimeStep(ts, dt_init));
267: PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, globXArray));
268: if (b_target >= global_batch_id && b_target < global_batch_id + ctx->batch_sz) PetscCall(PetscObjectSetName((PetscObject)globXArray[LAND_PACK_IDX(b_target % ctx->batch_sz, g_target)], "rho"));
269: // create fake particles
270: for (PetscInt b_id_0 = 0; b_id_0 < ctx->batch_sz; b_id_0 += numthreads) {
271: PetscReal *xx_t[LANDAU_MAX_GRIDS][MAX_NUM_THRDS], *yy_t[LANDAU_MAX_GRIDS][MAX_NUM_THRDS], *zz_t[LANDAU_MAX_GRIDS][MAX_NUM_THRDS], *wp_t[LANDAU_MAX_GRIDS][MAX_NUM_THRDS];
272: PetscInt Np_t[LANDAU_MAX_GRIDS][MAX_NUM_THRDS];
273: // make particles
274: for (int tid = 0; tid < numthreads; tid++) {
275: const PetscInt b_id = b_id_0 + tid;
276: if ((glb_b_id = global_batch_id + b_id) < NUserV) { // the ragged edge of the last batch
277: PetscInt Npp0 = a_Np + (glb_b_id % a_Np), NN; // fake user: number of particels in each dimension with add some load imbalance and diff (<2x)
278: for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
279: const PetscReal kT_m = ctx->k * ctx->thermal_temps[ctx->species_offset[grid]] / ctx->masses[ctx->species_offset[grid]] / (ctx->v_0 * ctx->v_0); /* theta = 2kT/mc^2 per species -- TODO */
280: ;
281: PetscReal lo[3] = {-ctx->radius[grid], -ctx->radius[grid], -ctx->radius[grid]}, hi[3] = {ctx->radius[grid], ctx->radius[grid], ctx->radius[grid]}, hp[3], vole; // would be nice to get box from DM
282: PetscInt Npi = Npp0, Npj = 2 * Npp0, Npk = 1;
283: if (dim == 2) lo[0] = 0; // Landau coordinate (r,z)
284: else Npi = Npj = Npk = Npp0;
285: // User: use glb_b_id to index into your data
286: NN = Npi * Npj * Npk; // make a regular grid of particles Npp x Npp
287: if (glb_b_id == b_target) {
288: nTargetP = NN;
289: PetscCall(PetscInfo(pack, "Target %" PetscInt_FMT " with %" PetscInt_FMT " particels\n", glb_b_id, NN));
290: }
291: Np_t[grid][tid] = NN;
292: PetscCall(PetscMalloc4(NN, &xx_t[grid][tid], NN, &yy_t[grid][tid], NN, &wp_t[grid][tid], dim == 2 ? 1 : NN, &zz_t[grid][tid]));
293: hp[0] = (hi[0] - lo[0]) / Npi;
294: hp[1] = (hi[1] - lo[1]) / Npj;
295: hp[2] = (hi[2] - lo[2]) / Npk;
296: if (dim == 2) hp[2] = 1;
297: PetscCall(PetscInfo(pack, " lo = %14.7e, hi = %14.7e; hp = %14.7e, %14.7e; kT_m = %g; \n", (double)lo[1], (double)hi[1], (double)hp[0], (double)hp[1], (double)kT_m)); // temp
298: vole = hp[0] * hp[1] * hp[2] * ctx->n[grid]; // fix for multi-species
299: PetscCall(PetscInfo(pack, "Vertex %" PetscInt_FMT ", grid %" PetscInt_FMT " with %" PetscInt_FMT " particles (diagnostic target = %" PetscInt_FMT ")\n", glb_b_id, grid, NN, b_target));
300: for (int pj = 0, pp = 0; pj < Npj; pj++) {
301: for (int pk = 0; pk < Npk; pk++) {
302: for (int pi = 0; pi < Npi; pi++, pp++) {
303: xx_t[grid][tid][pp] = lo[0] + hp[0] / 2.0 + pi * hp[0];
304: yy_t[grid][tid][pp] = lo[1] + hp[1] / 2.0 + pj * hp[1];
305: if (dim == 3) zz_t[grid][tid][pp] = lo[2] + hp[2] / 2.0 + pk * hp[2];
306: {
307: PetscReal x[] = {xx_t[grid][tid][pp], yy_t[grid][tid][pp], dim == 2 ? 0 : zz_t[grid][tid][pp]};
308: maxwellian(dim, x, kT_m, vole, &wp_t[grid][tid][pp]);
309: //PetscCall(PetscInfo(pack,"%" PetscInt_FMT ") x = %14.7e, %14.7e, %14.7e, n = %14.7e, w = %14.7e\n", pp, x[0], x[1], dim==2 ? 0 : x[2], ctx->n[grid], wp_t[grid][tid][pp])); // temp
310: if (glb_b_id == b_target) {
311: PetscReal v2 = 0, fact = dim == 2 ? 2.0 * PETSC_PI * x[0] : 1;
312: for (int i = 0; i < dim; ++i) v2 += PetscSqr(x[i]);
313: moments_0[0] += fact * wp_t[grid][tid][pp] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]];
314: moments_0[1] += fact * wp_t[grid][tid][pp] * ctx->n_0 * ctx->v_0 * ctx->masses[ctx->species_offset[grid]] * x[1]; // z-momentum
315: moments_0[2] += fact * wp_t[grid][tid][pp] * ctx->n_0 * ctx->v_0 * ctx->v_0 * ctx->masses[ctx->species_offset[grid]] * v2;
316: }
317: }
318: }
319: }
320: }
321: } // grid
322: } // active
323: } // fake threads
324: /* Create particle swarm */
325: PetscPragmaOMP(parallel for)
326: for (int tid = 0; tid < numthreads; tid++) {
327: const PetscInt b_id = b_id_0 + tid;
328: if ((glb_b_id = global_batch_id + b_id) < NUserV) { // the ragged edge of the last batch
329: //PetscCall(PetscInfo(pack,"Create swarms for 'glob' index %" PetscInt_FMT " create swarm\n",glb_b_id));
330: for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
331: PetscErrorCode ierr_t;
332: PetscSection section;
333: PetscInt Nf;
334: DM dm = grid_dm[grid];
335: ierr_t = DMGetLocalSection(dm, §ion);
336: ierr_t = PetscSectionGetNumFields(section, &Nf);
337: if (Nf != 1) ierr_t = 9999;
338: else {
339: ierr_t = DMViewFromOptions(dm, NULL, "-dm_view");
340: ierr_t = PetscInfo(pack, "call createSwarm [%" PetscInt_FMT ".%" PetscInt_FMT "] local batch index %" PetscInt_FMT "\n", b_id, grid, LAND_PACK_IDX(b_id, grid));
341: ierr_t = createSwarm(dm, dim, &globSwarmArray[LAND_PACK_IDX(b_id, grid)]);
342: }
343: if (ierr_t) ierr = ierr_t;
344: }
345: } // active
346: }
347: PetscCheck(ierr != 9999, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Only support one species per grid");
348: PetscCheck(!ierr, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Error in OMP loop. ierr = %d", (int)ierr);
349: // p --> g: make globMpArray & set X
350: PetscPragmaOMP(parallel for)
351: for (int tid = 0; tid < numthreads; tid++) {
352: const PetscInt b_id = b_id_0 + tid;
353: if ((glb_b_id = global_batch_id + b_id) < NUserV) {
354: for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
355: PetscErrorCode ierr_t;
356: DM dm = grid_dm[grid];
357: DM sw = globSwarmArray[LAND_PACK_IDX(b_id, grid)];
358: Vec subX = globXArray[LAND_PACK_IDX(b_id, grid)], work = t_fhat[grid][tid];
359: ierr_t = PetscInfo(pack, "particlesToGrid %" PetscInt_FMT ".%" PetscInt_FMT ") particlesToGrid for local batch %" PetscInt_FMT "\n", global_batch_id, grid, LAND_PACK_IDX(b_id, grid));
360: ierr_t = particlesToGrid(dm, sw, Np_t[grid][tid], tid, dim, xx_t[grid][tid], yy_t[grid][tid], zz_t[grid][tid], wp_t[grid][tid], subX, &globMpArray[LAND_PACK_IDX(b_id, grid)]);
361: if (ierr_t) ierr = ierr_t;
362: // u = M^_1 f_w
363: ierr_t = VecCopy(subX, work);
364: ierr_t = KSPSolve(t_ksp[grid][tid], work, subX);
365: if (ierr_t) ierr = ierr_t;
366: }
367: }
368: }
369: PetscCheck(!ierr, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Error in OMP loop. ierr = %d", (int)ierr);
370: /* Cleanup */
371: for (int tid = 0; tid < numthreads; tid++) {
372: const PetscInt b_id = b_id_0 + tid;
373: if ((glb_b_id = global_batch_id + b_id) < NUserV) {
374: PetscCall(PetscInfo(pack, "Free for global batch %" PetscInt_FMT " of %" PetscInt_FMT "\n", glb_b_id + 1, NUserV));
375: for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
376: PetscCall(PetscFree4(xx_t[grid][tid], yy_t[grid][tid], wp_t[grid][tid], zz_t[grid][tid]));
377: }
378: } // active
379: }
380: } // batches
381: if (b_target >= global_batch_id && b_target < global_batch_id + ctx->batch_sz) PetscCall(VecViewFromOptions(globXArray[LAND_PACK_IDX(b_target % ctx->batch_sz, g_target)], NULL, "-ex30_vec_view"));
382: PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, globXArray));
383: PetscCall(DMPlexLandauPrintNorms(X, 0));
384: // advance
385: PetscCall(TSSetSolution(ts, X));
386: PetscCall(PetscInfo(pack, "Advance vertex %" PetscInt_FMT " to %" PetscInt_FMT " (with padding)\n", global_batch_id, global_batch_id + ctx->batch_sz));
387: PetscCall(TSSolve(ts, X));
388: PetscCall(DMPlexLandauPrintNorms(X, 1));
389: PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, globXArray));
390: // map back to particles
391: for (PetscInt b_id_0 = 0; b_id_0 < ctx->batch_sz; b_id_0 += numthreads) {
392: PetscCall(PetscInfo(pack, "g2p: global batch %" PetscInt_FMT " of %" PetscInt_FMT ", Landau batch %" PetscInt_FMT " of %" PetscInt_FMT ": map back to particles\n", global_batch_id + 1, NUserV, b_id_0 + 1, ctx->batch_sz));
393: PetscPragmaOMP(parallel for)
394: for (int tid = 0; tid < numthreads; tid++) {
395: const PetscInt b_id = b_id_0 + tid;
396: if ((glb_b_id = global_batch_id + b_id) < NUserV) {
397: for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
398: PetscErrorCode ierr_t;
399: ierr_t = PetscInfo(pack, "gridToParticles: global batch %" PetscInt_FMT ", local batch b=%" PetscInt_FMT ", grid g=%" PetscInt_FMT ", index(b,g) %" PetscInt_FMT "\n", global_batch_id, b_id, grid, LAND_PACK_IDX(b_id, grid));
400: ierr_t = gridToParticles(grid_dm[grid], globSwarmArray[LAND_PACK_IDX(b_id, grid)], globXArray[LAND_PACK_IDX(b_id, grid)], t_fhat[grid][tid], globMpArray[LAND_PACK_IDX(b_id, grid)], g_Mass[grid]);
401: if (ierr_t) ierr = ierr_t;
402: }
403: }
404: }
405: PetscCheck(!ierr, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Error in OMP loop. ierr = %d", (int)ierr);
406: /* Cleanup, and get data */
407: PetscCall(PetscInfo(pack, "Cleanup batches %" PetscInt_FMT " to %" PetscInt_FMT "\n", b_id_0, b_id_0 + numthreads));
408: for (int tid = 0; tid < numthreads; tid++) {
409: const PetscInt b_id = b_id_0 + tid;
410: if ((glb_b_id = global_batch_id + b_id) < NUserV) {
411: for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
412: PetscDataType dtype;
413: PetscReal *wp, *coords;
414: DM sw = globSwarmArray[LAND_PACK_IDX(b_id, grid)];
415: PetscInt npoints, bs = 1;
416: PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wp)); // take data out here
417: if (glb_b_id == b_target) {
418: PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
419: PetscCall(DMSwarmGetLocalSize(sw, &npoints));
420: for (int p = 0; p < npoints; p++) {
421: PetscReal v2 = 0, fact = dim == 2 ? 2.0 * PETSC_PI * coords[p * dim + 0] : 1;
422: for (int i = 0; i < dim; ++i) v2 += PetscSqr(coords[p * dim + i]);
423: moments_1[0] += fact * wp[p] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]];
424: moments_1[1] += fact * wp[p] * ctx->n_0 * ctx->v_0 * ctx->masses[ctx->species_offset[grid]] * coords[p * dim + 1]; // z-momentum
425: moments_1[2] += fact * wp[p] * ctx->n_0 * ctx->v_0 * ctx->v_0 * ctx->masses[ctx->species_offset[grid]] * v2;
426: }
427: PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
428: }
429: PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wp));
430: PetscCall(DMDestroy(&globSwarmArray[LAND_PACK_IDX(b_id, grid)]));
431: PetscCall(MatDestroy(&globMpArray[LAND_PACK_IDX(b_id, grid)]));
432: }
433: }
434: }
435: } // thread batch
436: PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, globXArray));
437: } // user batch
438: /* Cleanup */
439: PetscCall(PetscFree(globXArray));
440: PetscCall(PetscFree(globSwarmArray));
441: PetscCall(PetscFree(globMpArray));
442: // clean up mass matrices
443: for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
444: PetscCall(MatDestroy(&g_Mass[grid]));
445: for (int tid = 0; tid < numthreads; tid++) {
446: PetscCall(VecDestroy(&t_fhat[grid][tid]));
447: PetscCall(KSPDestroy(&t_ksp[grid][tid]));
448: }
449: }
450: {
451: PetscReal relerr = PetscAbsReal((moments_1[2] - moments_0[2]) / moments_0[2]), logEerr = PetscLog10Real(relerr);
452: PetscCall(PetscInfo(X, "Total number density: %20.12e (%20.12e); x-momentum = %20.12e (%20.12e); energy = %20.12e (%20.12e) rerr = %e (log10(rerr) = %e, %" PetscInt_FMT "), %" 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_0[2], (double)relerr, (double)logEerr, -((-(PetscInt)logEerr + 1) / 2) * 2, nTargetP, numthreads));
453: PetscFunctionReturn(PETSC_SUCCESS);
454: }
455: PetscFunctionReturn(PETSC_SUCCESS);
456: }
458: int main(int argc, char **argv)
459: {
460: DM pack;
461: Vec X;
462: PetscInt dim = 2, nvert = 1, Np = 10, btarget = 0, gtarget = 0;
463: TS ts;
464: Mat J;
465: LandauCtx *ctx;
467: PetscFunctionBeginUser;
468: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
469: // process args
470: PetscOptionsBegin(PETSC_COMM_SELF, "", "Collision Options", "DMPLEX");
471: PetscCall(PetscOptionsInt("-number_spatial_vertices", "Number of user spatial vertices to be batched for Landau", "ex30.c", nvert, &nvert, NULL));
472: PetscCall(PetscOptionsInt("-dim", "Velocity space dimension", "ex30.c", dim, &dim, NULL));
473: PetscCall(PetscOptionsInt("-number_particles_per_dimension", "Number of particles per grid, with slight modification per spatial vertex, in each dimension of base Cartesian grid", "ex30.c", Np, &Np, NULL));
474: PetscCall(PetscOptionsInt("-view_vertex_target", "Batch to view with diagnostics", "ex30.c", btarget, &btarget, NULL));
475: PetscCheck(btarget < nvert, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Batch to view %" PetscInt_FMT " should be < number of vertices %" PetscInt_FMT, btarget, nvert);
476: PetscCall(PetscOptionsInt("-view_grid_target", "Grid to view with diagnostics", "ex30.c", gtarget, >arget, NULL));
477: PetscOptionsEnd();
478: /* Create a mesh */
479: PetscCall(DMPlexLandauCreateVelocitySpace(PETSC_COMM_SELF, dim, "", &X, &J, &pack));
480: PetscCall(DMSetUp(pack));
481: PetscCall(DMSetOutputSequenceNumber(pack, 0, 0.0));
482: PetscCall(DMGetApplicationContext(pack, &ctx));
483: PetscCheck(gtarget < ctx->num_grids, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Grid to view %" PetscInt_FMT " should be < number of grids %" PetscInt_FMT, gtarget, ctx->num_grids);
484: PetscCheck(nvert >= ctx->batch_sz, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number of vertices %" PetscInt_FMT " should be <= batch size %" PetscInt_FMT, nvert, ctx->batch_sz);
485: /* Create timestepping solver context */
486: PetscCall(TSCreate(PETSC_COMM_SELF, &ts));
487: PetscCall(TSSetDM(ts, pack));
488: PetscCall(TSSetIFunction(ts, NULL, DMPlexLandauIFunction, NULL));
489: PetscCall(TSSetIJacobian(ts, J, J, DMPlexLandauIJacobian, NULL));
490: PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
491: PetscCall(TSSetFromOptions(ts));
492: PetscCall(PetscObjectSetName((PetscObject)X, "X"));
493: // do particle advance
494: PetscCall(go(ts, X, nvert, Np, dim, btarget, gtarget));
495: PetscCall(MatZeroEntries(J)); // need to zero out so as to not reuse it in Landau's logic
496: /* clean up */
497: PetscCall(DMPlexLandauDestroyVelocitySpace(&pack));
498: PetscCall(TSDestroy(&ts));
499: PetscCall(VecDestroy(&X));
500: PetscCall(PetscFinalize());
501: return 0;
502: }
504: /*TEST
506: build:
507: requires: !complex p4est
509: testset:
510: requires: double defined(PETSC_USE_DMLANDAU_2D)
511: output_file: output/ex30_0.out
512: args: -dim 2 -petscspace_degree 3 -dm_landau_type p4est -dm_landau_num_species_grid 1,1,1 -dm_landau_amr_levels_max 0,0,0 \
513: -dm_landau_amr_post_refine 1 -number_particles_per_dimension 10 -dm_plex_hash_location \
514: -dm_landau_batch_size 2 -number_spatial_vertices 3 -dm_landau_batch_view_idx 1 -view_vertex_target 2 -view_grid_target 1 \
515: -dm_landau_n 1.000018,1,1e-6 -dm_landau_thermal_temps 2,1,1 -dm_landau_ion_masses 2,180 -dm_landau_ion_charges 1,18 \
516: -ftop_ksp_converged_reason -ftop_ksp_rtol 1e-10 -ftop_ksp_type lsqr -ftop_pc_type bjacobi -ftop_sub_pc_factor_shift_type nonzero -ftop_sub_pc_type lu \
517: -ksp_type preonly -pc_type lu -dm_landau_verbose 4 \
518: -ptof_ksp_type cg -ptof_pc_type jacobi -ptof_ksp_converged_reason -ptof_ksp_rtol 1e-12\
519: -snes_converged_reason -snes_monitor -snes_rtol 1e-14 -snes_stol 1e-14\
520: -ts_dt 0.01 -ts_rtol 1e-1 -ts_exact_final_time stepover -ts_max_snes_failures -1 -ts_max_steps 1 -ts_monitor -ts_type beuler -info :vec
522: test:
523: suffix: cpu
524: args: -dm_landau_device_type cpu
525: test:
526: suffix: kokkos
527: requires: kokkos_kernels
528: args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos
529: test:
530: suffix: cuda
531: requires: cuda
532: args: -dm_landau_device_type cuda -dm_mat_type aijcusparse -dm_vec_type cuda
534: testset:
535: requires: double !defined(PETSC_USE_DMLANDAU_2D)
536: output_file: output/ex30_3d.out
537: args: -dim 3 -petscspace_degree 2 -dm_landau_type p8est -dm_landau_num_species_grid 1,1,1 -dm_landau_amr_levels_max 0,0,0 \
538: -dm_landau_amr_post_refine 0 -number_particles_per_dimension 5 -dm_plex_hash_location \
539: -dm_landau_batch_size 1 -number_spatial_vertices 1 -dm_landau_batch_view_idx 0 -view_vertex_target 0 -view_grid_target 0 \
540: -dm_landau_n 1.000018,1,1e-6 -dm_landau_thermal_temps 2,1,1 -dm_landau_ion_masses 2,180 -dm_landau_ion_charges 1,18 \
541: -ftop_ksp_converged_reason -ftop_ksp_rtol 1e-12 -ftop_ksp_type cg -ftop_pc_type jacobi \
542: -ksp_type preonly -pc_type lu \
543: -ptof_ksp_type cg -ptof_pc_type jacobi -ptof_ksp_converged_reason -ptof_ksp_rtol 1e-6\
544: -snes_converged_reason -snes_monitor -snes_rtol 1e-9 -snes_stol 1e-9\
545: -ts_dt 0.1 -ts_exact_final_time stepover -ts_max_snes_failures -1 -ts_max_steps 1 -ts_monitor -ts_type beuler -info :vec
547: test:
548: suffix: cpu_3d
549: args: -dm_landau_device_type cpu
550: test:
551: suffix: kokkos_3d
552: requires: kokkos_kernels
553: args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos
554: test:
555: suffix: cuda_3d
556: requires: cuda
557: args: -dm_landau_device_type cuda -dm_mat_type aijcusparse -dm_vec_type cuda
559: TEST*/