1 static char help[] = "Grid based Landau collision operator with PIC interface with OpenMP setup. (one species per grid)\n";
2
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
10
11 */
12
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>
23
24 typedef struct {
25 Mat MpTrans;
26 Mat Mp;
27 Vec ff;
28 Vec uu;
29 } MatShellCtx;
30
31 typedef struct {
32 PetscInt v_target;
33 PetscInt g_target;
34 PetscInt global_vertex_id_0;
35 DM *globSwarmArray;
36 LandauCtx *ctx;
37 DM *grid_dm;
38 Mat *g_Mass;
39 Mat *globMpArray;
40 Vec *globXArray;
41 PetscBool print;
42 PetscBool print_entropy;
43 } PrintCtx;
44
MatMultMtM_SeqAIJ(Mat MtM,Vec xx,Vec yy)45 PetscErrorCode MatMultMtM_SeqAIJ(Mat MtM, Vec xx, Vec yy)
46 {
47 MatShellCtx *matshellctx;
48
49 PetscFunctionBeginUser;
50 PetscCall(MatShellGetContext(MtM, &matshellctx));
51 PetscCheck(matshellctx, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "No context");
52 PetscCall(MatMult(matshellctx->Mp, xx, matshellctx->ff));
53 PetscCall(MatMult(matshellctx->MpTrans, matshellctx->ff, yy));
54 PetscFunctionReturn(PETSC_SUCCESS);
55 }
56
MatMultAddMtM_SeqAIJ(Mat MtM,Vec xx,Vec yy,Vec zz)57 PetscErrorCode MatMultAddMtM_SeqAIJ(Mat MtM, Vec xx, Vec yy, Vec zz)
58 {
59 MatShellCtx *matshellctx;
60
61 PetscFunctionBeginUser;
62 PetscCall(MatShellGetContext(MtM, &matshellctx));
63 PetscCheck(matshellctx, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "No context");
64 PetscCall(MatMult(matshellctx->Mp, xx, matshellctx->ff));
65 PetscCall(MatMultAdd(matshellctx->MpTrans, matshellctx->ff, yy, zz));
66 PetscFunctionReturn(PETSC_SUCCESS);
67 }
68
createSwarm(const DM dm,PetscInt dim,DM * sw)69 PetscErrorCode createSwarm(const DM dm, PetscInt dim, DM *sw)
70 {
71 PetscInt Nc = 1;
72
73 PetscFunctionBeginUser;
74 PetscCall(DMCreate(PETSC_COMM_SELF, sw));
75 PetscCall(DMSetType(*sw, DMSWARM));
76 PetscCall(DMSetDimension(*sw, dim));
77 PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC));
78 PetscCall(DMSwarmSetCellDM(*sw, dm));
79 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", Nc, PETSC_REAL));
80 PetscCall(DMSwarmFinalizeFieldRegister(*sw));
81 PetscCall(DMSetFromOptions(*sw));
82 PetscCall(PetscObjectSetName((PetscObject)*sw, "Particle Grid"));
83 PetscFunctionReturn(PETSC_SUCCESS);
84 }
85
makeSwarm(DM sw,const PetscInt dim,const PetscInt Np,const PetscReal xx[],const PetscReal yy[],const PetscReal zz[])86 static PetscErrorCode makeSwarm(DM sw, const PetscInt dim, const PetscInt Np, const PetscReal xx[], const PetscReal yy[], const PetscReal zz[])
87 {
88 PetscReal *coords;
89 PetscDataType dtype;
90 PetscInt bs, p, zero = 0;
91
92 PetscFunctionBeginUser;
93 PetscCall(DMSwarmSetLocalSizes(sw, Np, zero));
94 PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
95 for (p = 0; p < Np; p++) {
96 coords[p * dim + 0] = xx[p];
97 coords[p * dim + 1] = yy[p];
98 if (dim == 3) coords[p * dim + 2] = zz[p];
99 }
100 PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
101 PetscCall(DMSwarmVectorDefineField(sw, "w_q"));
102 PetscFunctionReturn(PETSC_SUCCESS);
103 }
104
createMp(const DM dm,DM sw,Mat * Mp_out)105 static PetscErrorCode createMp(const DM dm, DM sw, Mat *Mp_out)
106 {
107 PetscBool removePoints = PETSC_TRUE;
108 Mat M_p;
109
110 PetscFunctionBeginUser;
111 // migrate after coords are set
112 PetscCall(DMSwarmMigrate(sw, removePoints));
113 //
114 PetscCall(PetscObjectSetName((PetscObject)sw, "Particle Grid"));
115
116 /* PetscInt N,*count,nmin=10000,nmax=0,ntot=0; */
117 /* // count */
118 /* PetscCall(DMSwarmCreatePointPerCellCount(sw, &N, &count)); */
119 /* for (int i=0, n; i< N ; i++) { */
120 /* if ((n=count[i]) > nmax) nmax = n; */
121 /* if (n < nmin) nmin = n; */
122 /* PetscCall(PetscInfo(dm, " %d) %d particles\n", i, n)); */
123 /* ntot += n; */
124 /* } */
125 /* PetscCall(PetscFree(count)); */
126 /* PetscCall(PetscInfo(dm, " %" PetscInt_FMT " max particle / cell, and %" PetscInt_FMT " min, ratio = %g, %" PetscInt_FMT " total\n", nmax, nmin, (double)nmax/(double)nmin,ntot)); */
127
128 /* This gives M f = \int_\Omega \phi f, which looks like a rhs for a PDE */
129 PetscCall(DMCreateMassMatrix(sw, dm, &M_p));
130 PetscCall(DMViewFromOptions(sw, NULL, "-ex30_sw_view"));
131 // output
132 *Mp_out = M_p;
133 PetscFunctionReturn(PETSC_SUCCESS);
134 }
135
particlesToGrid(const DM dm,DM sw,const PetscInt a_tid,const PetscInt dim,const PetscReal a_wp[],Vec rho,Mat M_p)136 static PetscErrorCode particlesToGrid(const DM dm, DM sw, const PetscInt a_tid, const PetscInt dim, const PetscReal a_wp[], Vec rho, Mat M_p)
137 {
138 PetscReal *wq;
139 PetscDataType dtype;
140 Vec ff;
141 PetscInt bs, p, Np;
142
143 PetscFunctionBeginUser;
144 PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wq));
145 PetscCall(DMSwarmGetLocalSize(sw, &Np));
146 for (p = 0; p < Np; p++) wq[p] = a_wp[p];
147 PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wq));
148 PetscCall(PetscObjectSetName((PetscObject)rho, "rho"));
149 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &ff));
150 PetscCall(PetscObjectSetName((PetscObject)ff, "weights"));
151 PetscCall(MatMultTranspose(M_p, ff, rho));
152 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &ff));
153 PetscFunctionReturn(PETSC_SUCCESS);
154 }
155
156 //
157 // add grid to arg 'sw.w_q'
158 //
gridToParticles(const DM dm,DM sw,const Vec rhs,Vec work_ferhs,Mat M_p,Mat Mass)159 PetscErrorCode gridToParticles(const DM dm, DM sw, const Vec rhs, Vec work_ferhs, Mat M_p, Mat Mass)
160 {
161 PetscBool is_lsqr;
162 KSP ksp;
163 Mat PM_p = NULL, MtM, D = NULL;
164 Vec ff;
165 PetscInt N, M, nzl;
166 MatShellCtx *matshellctx = NULL;
167 PC pc;
168
169 PetscFunctionBeginUser;
170 // 1) apply M in, for Moore-Penrose with mass: Mp (Mp' Mp)^-1 M
171 PetscCall(MatMult(Mass, rhs, work_ferhs));
172 // 2) pseudo-inverse, first part: (Mp' Mp)^-1
173 PetscCall(KSPCreate(PETSC_COMM_SELF, &ksp));
174 PetscCall(KSPSetType(ksp, KSPCG));
175 PetscCall(KSPGetPC(ksp, &pc));
176 PetscCall(PCSetType(pc, PCJACOBI));
177 PetscCall(KSPSetOptionsPrefix(ksp, "ftop_"));
178 PetscCall(KSPSetFromOptions(ksp));
179 PetscCall(PetscObjectTypeCompare((PetscObject)ksp, KSPLSQR, &is_lsqr));
180 if (!is_lsqr) {
181 PetscCall(MatGetLocalSize(M_p, &M, &N));
182 if (N > M) {
183 PetscCall(PetscInfo(ksp, " M (%" PetscInt_FMT ") < M (%" PetscInt_FMT ") more vertices than particles: revert to lsqr\n", M, N));
184 is_lsqr = PETSC_TRUE;
185 PetscCall(KSPSetType(ksp, KSPLSQR));
186 PetscCall(PCSetType(pc, PCNONE)); // should not happen, but could solve stable (Mp^T Mp), move projection Mp before solve
187 } else {
188 PetscCall(PetscNew(&matshellctx));
189 PetscCall(MatCreateVecs(M_p, &matshellctx->uu, &matshellctx->ff));
190 if (0) {
191 PetscCall(MatTransposeMatMult(M_p, M_p, MAT_INITIAL_MATRIX, 4, &MtM));
192 PetscCall(KSPSetOperators(ksp, MtM, MtM));
193 PetscCall(PetscInfo(M_p, "createMtM KSP with explicit Mp'Mp\n"));
194 PetscCall(MatViewFromOptions(MtM, NULL, "-ftop2_MtM_mat_view"));
195 } else {
196 PetscCall(MatCreateShell(PetscObjectComm((PetscObject)dm), N, N, PETSC_DECIDE, PETSC_DECIDE, matshellctx, &MtM));
197 PetscCall(MatTranspose(M_p, MAT_INITIAL_MATRIX, &matshellctx->MpTrans));
198 matshellctx->Mp = M_p;
199 PetscCall(MatShellSetOperation(MtM, MATOP_MULT, (PetscErrorCodeFn *)MatMultMtM_SeqAIJ));
200 PetscCall(MatShellSetOperation(MtM, MATOP_MULT_ADD, (PetscErrorCodeFn *)MatMultAddMtM_SeqAIJ));
201 PetscCall(MatCreateSeqAIJ(PETSC_COMM_SELF, N, N, 1, NULL, &D));
202 PetscCall(MatViewFromOptions(matshellctx->MpTrans, NULL, "-ftop2_MpT_mat_view"));
203 for (PetscInt i = 0; i < N; i++) {
204 const PetscScalar *vals;
205 const PetscInt *cols;
206 PetscScalar dot = 0;
207 PetscCall(MatGetRow(matshellctx->MpTrans, i, &nzl, &cols, &vals));
208 for (PetscInt ii = 0; ii < nzl; ii++) dot += PetscSqr(vals[ii]);
209 if (dot < PETSC_MACHINE_EPSILON) {
210 PetscCall(PetscInfo(ksp, "empty row in pseudo-inverse %d\n", (int)i));
211 is_lsqr = PETSC_TRUE; // empty rows
212 PetscCall(KSPSetType(ksp, KSPLSQR));
213 PetscCall(PCSetType(pc, PCNONE)); // should not happen, but could solve stable (Mp Mp^T), move projection Mp before solve
214 // clean up
215 PetscCall(MatDestroy(&matshellctx->MpTrans));
216 PetscCall(VecDestroy(&matshellctx->ff));
217 PetscCall(VecDestroy(&matshellctx->uu));
218 PetscCall(MatDestroy(&D));
219 PetscCall(MatDestroy(&MtM));
220 PetscCall(PetscFree(matshellctx));
221 D = NULL;
222 break;
223 }
224 PetscCall(MatSetValue(D, i, i, dot, INSERT_VALUES));
225 }
226 if (D) {
227 PetscCall(MatAssemblyBegin(D, MAT_FINAL_ASSEMBLY));
228 PetscCall(MatAssemblyEnd(D, MAT_FINAL_ASSEMBLY));
229 PetscCall(PetscInfo(M_p, "createMtMKSP Have %" PetscInt_FMT " eqs, nzl = %" PetscInt_FMT "\n", N, nzl));
230 PetscCall(KSPSetOperators(ksp, MtM, D));
231 PetscCall(MatViewFromOptions(D, NULL, "-ftop2_D_mat_view"));
232 PetscCall(MatViewFromOptions(M_p, NULL, "-ftop2_Mp_mat_view"));
233 PetscCall(MatViewFromOptions(matshellctx->MpTrans, NULL, "-ftop2_MpTranspose_mat_view"));
234 PetscCall(MatViewFromOptions(MtM, NULL, "-ftop2_MtM_mat_view"));
235 }
236 }
237 }
238 }
239 if (is_lsqr) {
240 PC pc2;
241 PetscBool is_bjac;
242 PetscCall(KSPGetPC(ksp, &pc2));
243 PetscCall(PetscObjectTypeCompare((PetscObject)pc2, PCBJACOBI, &is_bjac));
244 if (is_bjac) {
245 PetscCall(DMSwarmCreateMassMatrixSquare(sw, dm, &PM_p));
246 PetscCall(KSPSetOperators(ksp, M_p, PM_p));
247 } else {
248 PetscCall(KSPSetOperators(ksp, M_p, M_p));
249 }
250 PetscCall(MatViewFromOptions(M_p, NULL, "-ftop2_Mp_mat_view"));
251 }
252 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &ff)); // this grabs access
253 if (!is_lsqr) {
254 PetscCall(KSPSolve(ksp, work_ferhs, matshellctx->uu));
255 // 3) with Moore-Penrose apply Mp: M_p (Mp' Mp)^-1 M
256 PetscCall(MatMult(M_p, matshellctx->uu, ff));
257 if (D) PetscCall(MatDestroy(&D));
258 PetscCall(MatDestroy(&MtM));
259 if (matshellctx->MpTrans) PetscCall(MatDestroy(&matshellctx->MpTrans));
260 PetscCall(VecDestroy(&matshellctx->ff));
261 PetscCall(VecDestroy(&matshellctx->uu));
262 PetscCall(PetscFree(matshellctx));
263 } else {
264 // finally with LSQR apply M_p^\dagger
265 PetscCall(KSPSolveTranspose(ksp, work_ferhs, ff));
266 }
267 PetscCall(KSPDestroy(&ksp));
268 PetscCall(MatDestroy(&PM_p));
269 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &ff));
270 PetscFunctionReturn(PETSC_SUCCESS);
271 }
272
273 #define EX30_MAX_NUM_THRDS 12
274 #define EX30_MAX_BATCH_SZ 1024
275 //
276 // add grid to arg 'globSwarmArray[].w_q'
277 //
gridToParticles_private(DM grid_dm[],DM globSwarmArray[],const PetscInt dim,const PetscInt v_target,const PetscInt numthreads,const PetscInt num_vertices,const PetscInt global_vertex_id,Mat globMpArray[],Mat g_Mass[],Vec t_fhat[][EX30_MAX_NUM_THRDS],PetscReal moments[],Vec globXArray[],LandauCtx * ctx)278 PetscErrorCode gridToParticles_private(DM grid_dm[], DM globSwarmArray[], const PetscInt dim, const PetscInt v_target, const PetscInt numthreads, const PetscInt num_vertices, const PetscInt global_vertex_id, Mat globMpArray[], Mat g_Mass[], Vec t_fhat[][EX30_MAX_NUM_THRDS], PetscReal moments[], Vec globXArray[], LandauCtx *ctx)
279 {
280 PetscErrorCode ierr = (PetscErrorCode)0; // used for inside thread loops
281
282 PetscFunctionBeginUser;
283 // map back to particles
284 for (PetscInt v_id_0 = 0; v_id_0 < ctx->batch_sz; v_id_0 += numthreads) {
285 PetscCall(PetscInfo(grid_dm[0], "g2p: global batch %" PetscInt_FMT " of %" PetscInt_FMT ", Landau batch %" PetscInt_FMT " of %" PetscInt_FMT ": map back to particles\n", global_vertex_id + 1, num_vertices, v_id_0 + 1, ctx->batch_sz));
286 //PetscPragmaOMP(parallel for)
287 for (PetscInt tid = 0; tid < numthreads; tid++) {
288 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id + v_id;
289 if (glb_v_id < num_vertices) {
290 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
291 PetscErrorCode ierr_t;
292 ierr_t = PetscInfo(grid_dm[0], "gridToParticles: global batch %" PetscInt_FMT ", local batch b=%" PetscInt_FMT ", grid g=%" PetscInt_FMT ", index(b,g) %" PetscInt_FMT "\n", global_vertex_id, v_id, grid, LAND_PACK_IDX(v_id, grid));
293 ierr_t = gridToParticles(grid_dm[grid], globSwarmArray[LAND_PACK_IDX(v_id, grid)], globXArray[LAND_PACK_IDX(v_id, grid)], t_fhat[grid][tid], globMpArray[LAND_PACK_IDX(v_id, grid)], g_Mass[grid]);
294 if (ierr_t) ierr = ierr_t;
295 }
296 }
297 }
298 PetscCheck(!ierr, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Error in OMP loop. ierr = %d", (int)ierr);
299 /* Get moments */
300 PetscCall(PetscInfo(grid_dm[0], "Cleanup batches %" PetscInt_FMT " to %" PetscInt_FMT "\n", v_id_0, v_id_0 + numthreads));
301 for (PetscInt tid = 0; tid < numthreads; tid++) {
302 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id + v_id;
303 if (glb_v_id == v_target) {
304 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
305 PetscDataType dtype;
306 PetscReal *wp, *coords;
307 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
308 PetscInt npoints, bs = 1;
309 PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wp)); // take data out here
310 PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
311 PetscCall(DMSwarmGetLocalSize(sw, &npoints));
312 for (PetscInt p = 0; p < npoints; p++) {
313 PetscReal v2 = 0, fact = (dim == 2) ? 2.0 * PETSC_PI * coords[p * dim + 0] : 1, w = fact * wp[p] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]];
314 for (PetscInt i = 0; i < dim; ++i) v2 += PetscSqr(coords[p * dim + i]);
315 moments[0] += w;
316 moments[1] += w * ctx->v_0 * coords[p * dim + 1]; // z-momentum
317 moments[2] += w * 0.5 * ctx->v_0 * ctx->v_0 * v2;
318 }
319 PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wp));
320 PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
321 }
322 const PetscReal N_inv = 1 / moments[0];
323 PetscCall(PetscInfo(grid_dm[0], "gridToParticles_private [%" PetscInt_FMT "], n = %g\n", v_id, (double)moments[0]));
324 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
325 PetscDataType dtype;
326 PetscReal *wp, *coords;
327 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
328 PetscInt npoints, bs = 1;
329 PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wp)); // take data out here
330 PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
331 PetscCall(DMSwarmGetLocalSize(sw, &npoints));
332 for (PetscInt p = 0; p < npoints; p++) {
333 const PetscReal fact = dim == 2 ? 2.0 * PETSC_PI * coords[p * dim + 0] : 1, w = fact * wp[p] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]], ww = w * N_inv;
334 if (w > PETSC_REAL_MIN) {
335 moments[3] -= ww * PetscLogReal(ww);
336 PetscCheck(ww < 1 - PETSC_MACHINE_EPSILON, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "ww (%g) > 1", (double)ww);
337 } else moments[4] -= w; // keep track of density that is lost
338 }
339 PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wp));
340 PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
341 }
342 }
343 } // thread batch
344 } // batch
345 PetscFunctionReturn(PETSC_SUCCESS);
346 }
347
maxwellian(PetscInt dim,const PetscReal x[],PetscReal kt_m,PetscReal n,PetscReal shift,PetscScalar * u)348 static void maxwellian(PetscInt dim, const PetscReal x[], PetscReal kt_m, PetscReal n, PetscReal shift, PetscScalar *u)
349 {
350 PetscInt i;
351 PetscReal v2 = 0, theta = 2.0 * kt_m; /* theta = 2kT/mc^2 */
352
353 if (shift != 0.) {
354 v2 = 0;
355 for (i = 0; i < dim - 1; ++i) v2 += x[i] * x[i];
356 v2 += (x[dim - 1] - shift) * (x[dim - 1] - shift);
357 /* evaluate the shifted Maxwellian */
358 u[0] += n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
359 } else {
360 /* compute the exponents, v^2 */
361 for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
362 /* evaluate the Maxwellian */
363 u[0] += n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
364 }
365 }
366
PostStep(TS ts)367 static PetscErrorCode PostStep(TS ts)
368 {
369 PetscInt n, dim, nDMs, v_id;
370 PetscReal t;
371 LandauCtx *ctx;
372 Vec X;
373 PrintCtx *printCtx;
374 DM pack;
375 PetscReal moments[5], e_grid[LANDAU_MAX_GRIDS];
376
377 PetscFunctionBeginUser;
378 PetscCall(TSGetApplicationContext(ts, &printCtx));
379 if (!printCtx->print && !printCtx->print_entropy) PetscFunctionReturn(PETSC_SUCCESS);
380 ctx = printCtx->ctx;
381 if (printCtx->v_target < printCtx->global_vertex_id_0 || printCtx->v_target >= printCtx->global_vertex_id_0 + ctx->batch_sz) PetscFunctionReturn(PETSC_SUCCESS);
382 for (PetscInt i = 0; i < 5; i++) moments[i] = 0;
383 for (PetscInt i = 0; i < LANDAU_MAX_GRIDS; i++) e_grid[i] = 0;
384 v_id = printCtx->v_target % ctx->batch_sz;
385 PetscCall(TSGetDM(ts, &pack));
386 PetscCall(DMGetDimension(pack, &dim));
387 PetscCall(DMCompositeGetNumberDM(pack, &nDMs)); // number of vertices * number of grids
388 PetscCall(TSGetSolution(ts, &X));
389 PetscCall(TSGetStepNumber(ts, &n));
390 PetscCall(TSGetTime(ts, &t));
391 PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, printCtx->globXArray));
392 if (printCtx->print_entropy && printCtx->v_target >= 0 && 0) {
393 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
394 PetscDataType dtype;
395 PetscReal *wp, *coords;
396 DM sw = printCtx->globSwarmArray[LAND_PACK_IDX(v_id, grid)];
397 Vec work, subX = printCtx->globXArray[LAND_PACK_IDX(v_id, grid)];
398 PetscInt bs, NN;
399 // C-G moments
400 PetscCall(VecDuplicate(subX, &work));
401 PetscCall(gridToParticles(printCtx->grid_dm[grid], sw, subX, work, printCtx->globMpArray[LAND_PACK_IDX(v_id, grid)], printCtx->g_Mass[grid]));
402 PetscCall(VecDestroy(&work));
403 // moments
404 PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
405 PetscCall(DMSwarmGetLocalSize(sw, &NN));
406 PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wp));
407 for (PetscInt pp = 0; pp < NN; pp++) {
408 PetscReal v2 = 0, fact = (dim == 2) ? 2.0 * PETSC_PI * coords[pp * dim + 0] : 1, w = fact * wp[pp] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]];
409 for (PetscInt i = 0; i < dim; ++i) v2 += PetscSqr(coords[pp * dim + i]);
410 moments[0] += w;
411 moments[1] += w * ctx->v_0 * coords[pp * dim + 1]; // z-momentum
412 moments[2] += w * 0.5 * ctx->v_0 * ctx->v_0 * v2;
413 e_grid[grid] += w * 0.5 * ctx->v_0 * ctx->v_0 * v2;
414 }
415 PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
416 PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wp));
417 }
418 // entropy
419 const PetscReal N_inv = 1 / moments[0];
420 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
421 PetscDataType dtype;
422 PetscReal *wp, *coords;
423 DM sw = printCtx->globSwarmArray[LAND_PACK_IDX(v_id, grid)];
424 PetscInt bs, NN;
425 PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
426 PetscCall(DMSwarmGetLocalSize(sw, &NN));
427 PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wp));
428 for (PetscInt pp = 0; pp < NN; pp++) {
429 PetscReal fact = (dim == 2) ? 2.0 * PETSC_PI * coords[pp * dim + 0] : 1, w = fact * wp[pp] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]], ww = w * N_inv;
430 if (w > PETSC_REAL_MIN) {
431 moments[3] -= ww * PetscLogReal(ww);
432 } else moments[4] -= w;
433 }
434 PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
435 PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wp));
436 }
437 PetscCall(PetscInfo(X, "%4d) time %e, Landau particle moments: 0: %18.12e 1: %19.12e 2: %18.12e entropy: %e loss %e. energy = %e + %e + %e\n", (int)n, (double)t, (double)moments[0], (double)moments[1], (double)moments[2], (double)moments[3], (double)(moments[4] / moments[0]), (double)e_grid[0], (double)e_grid[1], (double)e_grid[2]));
438 }
439 if (printCtx->print && printCtx->g_target >= 0) {
440 PetscInt grid = printCtx->g_target, id;
441 static PetscReal last_t = -100000, period = .5;
442 if (last_t == -100000) last_t = -period + t;
443 if (t >= last_t + period) {
444 last_t = t;
445 PetscCall(DMGetOutputSequenceNumber(ctx->plex[grid], &id, NULL));
446 PetscCall(DMSetOutputSequenceNumber(ctx->plex[grid], id + 1, t));
447 PetscCall(VecViewFromOptions(printCtx->globXArray[LAND_PACK_IDX(v_id % ctx->batch_sz, grid)], NULL, "-ex30_vec_view"));
448 if (ctx->num_grids > grid + 1) {
449 PetscCall(DMSetOutputSequenceNumber(ctx->plex[grid + 1], id + 1, t));
450 PetscCall(VecViewFromOptions(printCtx->globXArray[LAND_PACK_IDX(v_id % ctx->batch_sz, grid + 1)], NULL, "-ex30_vec_view2"));
451 }
452 PetscCall(PetscInfo(X, "%4d) time %e View\n", (int)n, (double)t));
453 }
454 }
455 PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, printCtx->globXArray));
456 PetscFunctionReturn(PETSC_SUCCESS);
457 }
458
go(TS ts,Vec X,const PetscInt num_vertices,const PetscInt a_Np,const PetscInt dim,const PetscInt v_target,const PetscInt g_target,PetscReal shift,PetscBool use_uniform_particle_grid)459 PetscErrorCode go(TS ts, Vec X, const PetscInt num_vertices, const PetscInt a_Np, const PetscInt dim, const PetscInt v_target, const PetscInt g_target, PetscReal shift, PetscBool use_uniform_particle_grid)
460 {
461 DM pack, *globSwarmArray, grid_dm[LANDAU_MAX_GRIDS];
462 Mat *globMpArray, g_Mass[LANDAU_MAX_GRIDS];
463 KSP t_ksp[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS];
464 Vec t_fhat[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS];
465 PetscInt nDMs;
466 PetscErrorCode ierr = (PetscErrorCode)0; // used for inside thread loops
467 #if defined(PETSC_HAVE_OPENMP) && defined(PETSC_HAVE_THREADSAFETY)
468 PetscInt numthreads = PetscNumOMPThreads;
469 #else
470 PetscInt numthreads = 1;
471 #endif
472 LandauCtx *ctx;
473 Vec *globXArray;
474 PetscReal moments_0[5], moments_1a[5], moments_1b[5], dt_init;
475 PrintCtx *printCtx;
476
477 PetscFunctionBeginUser;
478 PetscCheck(numthreads <= EX30_MAX_NUM_THRDS, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Too many threads %" PetscInt_FMT " > %d", numthreads, EX30_MAX_NUM_THRDS);
479 PetscCheck(numthreads > 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number threads %" PetscInt_FMT " > %d", numthreads, EX30_MAX_NUM_THRDS);
480 PetscCall(TSGetDM(ts, &pack));
481 PetscCall(DMGetApplicationContext(pack, &ctx));
482 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);
483 PetscCall(DMCompositeGetNumberDM(pack, &nDMs)); // number of vertices * number of grids
484 PetscCall(PetscInfo(pack, "Have %" PetscInt_FMT " total grids, with %" PetscInt_FMT " Landau local batched and %" PetscInt_FMT " global items (vertices) %d DMs\n", ctx->num_grids, ctx->batch_sz, num_vertices, (int)nDMs));
485 PetscCall(PetscMalloc(sizeof(*globXArray) * nDMs, &globXArray));
486 PetscCall(PetscMalloc(sizeof(*globMpArray) * nDMs, &globMpArray));
487 PetscCall(PetscMalloc(sizeof(*globSwarmArray) * nDMs, &globSwarmArray));
488 // print ctx
489 PetscCall(PetscNew(&printCtx));
490 PetscCall(TSSetApplicationContext(ts, printCtx));
491 printCtx->v_target = v_target;
492 printCtx->g_target = g_target;
493 printCtx->ctx = ctx;
494 printCtx->globSwarmArray = globSwarmArray;
495 printCtx->grid_dm = grid_dm;
496 printCtx->globMpArray = globMpArray;
497 printCtx->g_Mass = g_Mass;
498 printCtx->globXArray = globXArray;
499 printCtx->print_entropy = PETSC_FALSE;
500 PetscOptionsBegin(PETSC_COMM_SELF, "", "Print Options", "DMPLEX");
501 PetscCall(PetscOptionsBool("-print_entropy", "Print entropy and moments at each time step", "ex30.c", printCtx->print_entropy, &printCtx->print_entropy, NULL));
502 PetscOptionsEnd();
503 // view
504 PetscCall(DMViewFromOptions(ctx->plex[g_target], NULL, "-ex30_dm_view"));
505 if (ctx->num_grids > g_target + 1) PetscCall(DMViewFromOptions(ctx->plex[g_target + 1], NULL, "-ex30_dm_view2"));
506 // create mesh mass matrices
507 PetscCall(VecZeroEntries(X));
508 PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, globXArray)); // just to duplicate
509 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
510 Vec subX = globXArray[LAND_PACK_IDX(0, grid)];
511 DM dm = ctx->plex[grid];
512 PetscSection s;
513 grid_dm[grid] = dm;
514 PetscCall(DMCreateMassMatrix(dm, dm, &g_Mass[grid]));
515 //
516 PetscCall(DMGetLocalSection(dm, &s));
517 PetscCall(DMPlexCreateClosureIndex(dm, s));
518 for (PetscInt tid = 0; tid < numthreads; tid++) {
519 PC pc;
520 PetscCall(VecDuplicate(subX, &t_fhat[grid][tid]));
521 PetscCall(KSPCreate(PETSC_COMM_SELF, &t_ksp[grid][tid]));
522 PetscCall(KSPSetType(t_ksp[grid][tid], KSPCG));
523 PetscCall(KSPGetPC(t_ksp[grid][tid], &pc));
524 PetscCall(PCSetType(pc, PCJACOBI));
525 PetscCall(KSPSetOptionsPrefix(t_ksp[grid][tid], "ptof_"));
526 PetscCall(KSPSetOperators(t_ksp[grid][tid], g_Mass[grid], g_Mass[grid]));
527 PetscCall(KSPSetFromOptions(t_ksp[grid][tid]));
528 }
529 }
530 PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, globXArray));
531 PetscCall(TSGetTimeStep(ts, &dt_init)); // we could have an adaptive time stepper
532 // loop over all vertices in chucks that are batched for TSSolve
533 for (PetscInt i = 0; i < 5; i++) moments_0[i] = moments_1a[i] = moments_1b[i] = 0;
534 for (PetscInt global_vertex_id_0 = 0; global_vertex_id_0 < num_vertices; global_vertex_id_0 += ctx->batch_sz, shift /= 2) { // outer vertex loop
535 PetscCall(TSSetTime(ts, 0));
536 PetscCall(TSSetStepNumber(ts, 0));
537 PetscCall(TSSetTimeStep(ts, dt_init));
538 PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, globXArray));
539 printCtx->global_vertex_id_0 = global_vertex_id_0;
540 if (v_target >= global_vertex_id_0 && v_target < global_vertex_id_0 + ctx->batch_sz) {
541 PetscCall(PetscObjectSetName((PetscObject)globXArray[LAND_PACK_IDX(v_target % ctx->batch_sz, g_target)], "rho"));
542 printCtx->print = PETSC_TRUE;
543 } else printCtx->print = PETSC_FALSE;
544 // create fake particles in batches with threads
545 for (PetscInt v_id_0 = 0; v_id_0 < ctx->batch_sz; v_id_0 += numthreads) {
546 PetscReal *xx_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS], *yy_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS], *zz_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS], *wp_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS] /* , radiuses[80000] */;
547 PetscInt Np_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS];
548 // make particles
549 for (PetscInt tid = 0; tid < numthreads; tid++) {
550 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
551 if (glb_v_id < num_vertices) { // the ragged edge (in last batch)
552 PetscInt Npp0 = a_Np + (glb_v_id % (a_Np / 10 + 1)), nTargetP[LANDAU_MAX_GRIDS]; // n of particels in each dim with load imbalance
553 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
554 // for (PetscInt sp = ctx->species_offset[grid], i0 = 0; sp < ctx->species_offset[grid + 1]; sp++, i0++) {
555 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 */
556 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
557 PetscInt Npi = Npp0, Npj = 2 * Npp0, Npk = 1;
558 PetscRandom rand;
559 PetscReal sigma = ctx->thermal_speed[grid] / ctx->thermal_speed[0], p2_shift = grid == 0 ? shift : -shift; // symmetric shift of e vs ions
560 PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &rand));
561 PetscCall(PetscRandomSetInterval(rand, 0., 1.));
562 PetscCall(PetscRandomSetFromOptions(rand));
563 if (dim == 2) lo[0] = 0; // Landau coordinate (r,z)
564 else Npi = Npj = Npk = Npp0;
565 // User: use glb_v_id to index into your data
566 const PetscInt NNreal = Npi * Npj * Npk, NN = NNreal + (dim == 2 ? 3 : 6); // make room for bounding box
567 Np_t[grid][tid] = NN;
568 if (glb_v_id == v_target) nTargetP[grid] = NN;
569 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]));
570 hp[0] = (hi[0] - lo[0]) / Npi;
571 hp[1] = (hi[1] - lo[1]) / Npj;
572 hp[2] = (hi[2] - lo[2]) / Npk;
573 if (dim == 2) hp[2] = 1;
574 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
575 vole = hp[0] * hp[1] * hp[2] * ctx->n[grid]; // fix for multi-species
576 PetscCall(PetscInfo(pack, "Vertex %" PetscInt_FMT ", grid %" PetscInt_FMT " with %" PetscInt_FMT " particles (diagnostic target = %" PetscInt_FMT ")\n", glb_v_id, grid, NN, v_target));
577 for (PetscInt pj = 0, pp = 0; pj < Npj; pj++) {
578 for (PetscInt pk = 0; pk < Npk; pk++) {
579 for (PetscInt pi = 0; pi < Npi; pi++, pp++) {
580 PetscReal p_shift = p2_shift;
581 wp_t[grid][tid][pp] = 0;
582 if (use_uniform_particle_grid) {
583 xx_t[grid][tid][pp] = lo[0] + hp[0] / 2.0 + pi * hp[0];
584 yy_t[grid][tid][pp] = lo[1] + hp[1] / 2.0 + pj * hp[1];
585 if (dim == 3) zz_t[grid][tid][pp] = lo[2] + hp[2] / 2.0 + pk * hp[2];
586 PetscReal x[] = {xx_t[grid][tid][pp], yy_t[grid][tid][pp], dim == 2 ? 0 : zz_t[grid][tid][pp]};
587 p_shift *= ctx->thermal_speed[grid] / ctx->v_0;
588 if (ctx->sphere && PetscSqrtReal(PetscSqr(xx_t[grid][tid][pp]) + PetscSqr(yy_t[grid][tid][pp])) > 0.92 * hi[0]) {
589 wp_t[grid][tid][pp] = 0;
590 } else {
591 maxwellian(dim, x, kT_m, vole, p_shift, &wp_t[grid][tid][pp]);
592 if (ctx->num_grids == 1 && shift != 0) { // bi-maxwellian, electron plasma
593 maxwellian(dim, x, kT_m, vole, -p_shift, &wp_t[grid][tid][pp]); // symmetric shift of electron plasma
594 }
595 }
596 } else {
597 PetscReal u1, u2;
598 do {
599 do {
600 PetscCall(PetscRandomGetValueReal(rand, &u1));
601 } while (u1 == 0);
602 PetscCall(PetscRandomGetValueReal(rand, &u2));
603 //compute z0 and z1
604 PetscReal mag = sigma * PetscSqrtReal(-2.0 * PetscLogReal(u1)); // is this the same scale grid Maxwellian? t_therm = sigma
605 xx_t[grid][tid][pp] = mag * PetscCosReal(2.0 * PETSC_PI * u2);
606 yy_t[grid][tid][pp] = mag * PetscSinReal(2.0 * PETSC_PI * u2);
607 if (dim == 2 && xx_t[grid][tid][pp] < lo[0]) xx_t[grid][tid][pp] = -xx_t[grid][tid][pp];
608 if (dim == 3) zz_t[grid][tid][pp] = lo[2] + hp[2] / 2.0 + pk * hp[2];
609 if (!ctx->sphere) {
610 if (dim == 2 && xx_t[grid][tid][pp] < 0) xx_t[grid][tid][pp] = -xx_t[grid][tid][pp]; // ???
611 else if (dim == 3) {
612 while (zz_t[grid][tid][pp] >= hi[2] || zz_t[grid][tid][pp] <= lo[2]) zz_t[grid][tid][pp] *= .9;
613 }
614 while (xx_t[grid][tid][pp] >= hi[0] || xx_t[grid][tid][pp] <= lo[0]) xx_t[grid][tid][pp] *= .9;
615 while (yy_t[grid][tid][pp] >= hi[1] || yy_t[grid][tid][pp] <= lo[1]) yy_t[grid][tid][pp] *= .9;
616 } else { // 2D
617 //if (glb_v_id == v_target && pp < 80000) radiuses[pp] = PetscSqrtReal(PetscSqr(xx_t[grid][tid][pp]) + PetscSqr(yy_t[grid][tid][pp]));
618 while (PetscSqrtReal(PetscSqr(xx_t[grid][tid][pp]) + PetscSqr(yy_t[grid][tid][pp])) > 0.92 * hi[0]) { // safety factor for facets of sphere
619 xx_t[grid][tid][pp] *= .9;
620 yy_t[grid][tid][pp] *= .9;
621 }
622 }
623 if (ctx->num_grids == 1 && pp % 2 == 0) p_shift = 0; // one species, split bi-max
624 p_shift *= ctx->thermal_speed[grid] / ctx->v_0;
625 if (dim == 3) zz_t[grid][tid][pp] += p_shift;
626 else yy_t[grid][tid][pp] += p_shift;
627 wp_t[grid][tid][pp] += ctx->n[grid] / NNreal * PetscSqrtReal(ctx->masses[ctx->species_offset[grid]] / ctx->masses[0]);
628 if (p_shift <= 0) break; // add bi-max for electron plasma only
629 p_shift = -p_shift;
630 } while (ctx->num_grids == 1); // add bi-max for electron plasma only
631 }
632 {
633 if (glb_v_id == v_target) {
634 PetscReal x[] = {xx_t[grid][tid][pp], yy_t[grid][tid][pp], dim == 2 ? 0 : zz_t[grid][tid][pp]};
635 PetscReal v2 = 0, fact = dim == 2 ? 2.0 * PETSC_PI * x[0] : 1, w = fact * wp_t[grid][tid][pp] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]];
636 for (PetscInt i = 0; i < dim; ++i) v2 += PetscSqr(x[i]);
637 moments_0[0] += w; // not thread safe
638 moments_0[1] += w * ctx->v_0 * x[1]; // z-momentum
639 moments_0[2] += w * 0.5 * ctx->v_0 * ctx->v_0 * v2;
640 }
641 }
642 }
643 }
644 }
645 if (dim == 2) { // fix bounding box
646 PetscInt pp = NNreal;
647 wp_t[grid][tid][pp] = 0;
648 xx_t[grid][tid][pp] = 1.e-7;
649 yy_t[grid][tid][pp++] = hi[1] - 5.e-7;
650 wp_t[grid][tid][pp] = 0;
651 xx_t[grid][tid][pp] = hi[0] - 5.e-7;
652 yy_t[grid][tid][pp++] = 0;
653 wp_t[grid][tid][pp] = 0;
654 xx_t[grid][tid][pp] = 1.e-7;
655 yy_t[grid][tid][pp++] = lo[1] + 5.e-7;
656 } else {
657 const PetscInt p0 = NNreal;
658 for (PetscInt pj = 0; pj < 6; pj++) xx_t[grid][tid][p0 + pj] = yy_t[grid][tid][p0 + pj] = zz_t[grid][tid][p0 + pj] = wp_t[grid][tid][p0 + pj] = 0;
659 xx_t[grid][tid][p0 + 0] = lo[0];
660 xx_t[grid][tid][p0 + 1] = hi[0];
661 yy_t[grid][tid][p0 + 2] = lo[1];
662 yy_t[grid][tid][p0 + 3] = hi[1];
663 zz_t[grid][tid][p0 + 4] = lo[2];
664 zz_t[grid][tid][p0 + 5] = hi[2];
665 }
666 PetscCall(PetscRandomDestroy(&rand));
667 }
668 // entropy init, need global n
669 if (glb_v_id == v_target) {
670 const PetscReal N_inv = 1 / moments_0[0];
671 PetscCall(PetscInfo(pack, "Target %" PetscInt_FMT " with %" PetscInt_FMT " particels\n", glb_v_id, nTargetP[0]));
672 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
673 const PetscInt NN = nTargetP[grid];
674 for (PetscInt pp = 0; pp < NN; pp++) {
675 const PetscReal fact = dim == 2 ? 2.0 * PETSC_PI * xx_t[grid][tid][pp] : 1, w = fact * ctx->n_0 * ctx->masses[ctx->species_offset[grid]] * wp_t[grid][tid][pp], ww = w * N_inv;
676 if (w > PETSC_REAL_MIN) {
677 moments_0[3] -= ww * PetscLogReal(ww);
678 PetscCheck(ww < 1 - PETSC_MACHINE_EPSILON, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "ww (%g) > 1", (double)ww);
679 } else moments_0[4] -= w;
680 }
681 } // grid
682 } // target
683 } // active
684 } // threads
685 /* Create particle swarm */
686 for (PetscInt tid = 0; tid < numthreads; tid++) {
687 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
688 if (glb_v_id < num_vertices) { // the ragged edge of the last batch
689 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
690 PetscSection section;
691 PetscInt Nf;
692 DM dm = grid_dm[grid];
693 PetscCall(DMGetLocalSection(dm, §ion));
694 PetscCall(PetscSectionGetNumFields(section, &Nf));
695 PetscCheck(Nf == 1, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Only one species per grid supported -- todo");
696 PetscCall(DMViewFromOptions(dm, NULL, "-dm_view"));
697 PetscCall(PetscInfo(pack, "call createSwarm [%" PetscInt_FMT ".%" PetscInt_FMT "] local block index %" PetscInt_FMT "\n", v_id, grid, LAND_PACK_IDX(v_id, grid)));
698 PetscCall(createSwarm(dm, dim, &globSwarmArray[LAND_PACK_IDX(v_id, grid)]));
699 }
700 } // active
701 } // threads
702 PetscCheck(ierr != 9999, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Only support one species per grid");
703 // make globMpArray
704 PetscPragmaOMP(parallel for)
705 for (PetscInt tid = 0; tid < numthreads; tid++) {
706 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
707 if (glb_v_id < num_vertices) {
708 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
709 // for (PetscInt sp = ctx->species_offset[grid], i0 = 0; sp < ctx->species_offset[grid + 1]; sp++, i0++) -- loop over species for Nf > 1 -- TODO
710 PetscErrorCode ierr_t;
711 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
712 ierr_t = PetscInfo(pack, "makeSwarm %" PetscInt_FMT ".%" PetscInt_FMT ") for block %" PetscInt_FMT "\n", v_id, grid, LAND_PACK_IDX(v_id, grid));
713 ierr_t = makeSwarm(sw, dim, Np_t[grid][tid], xx_t[grid][tid], yy_t[grid][tid], zz_t[grid][tid]);
714 if (ierr_t) ierr = ierr_t;
715 }
716 }
717 }
718 for (PetscInt tid = 0; tid < numthreads; tid++) {
719 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
720 if (glb_v_id < num_vertices) {
721 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
722 DM dm = grid_dm[grid];
723 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
724 PetscCall(PetscInfo(pack, "createMp %" PetscInt_FMT ".%" PetscInt_FMT ") for block %" PetscInt_FMT "\n", v_id, grid, LAND_PACK_IDX(v_id, grid)));
725 PetscCall(createMp(dm, sw, &globMpArray[LAND_PACK_IDX(v_id, grid)]));
726 PetscCall(MatViewFromOptions(globMpArray[LAND_PACK_IDX(v_id, grid)], NULL, "-mp_mat_view"));
727 }
728 }
729 }
730 // p --> g: set X
731 // PetscPragmaOMP(parallel for)
732 for (PetscInt tid = 0; tid < numthreads; tid++) {
733 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
734 if (glb_v_id < num_vertices) {
735 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
736 PetscErrorCode ierr_t;
737 DM dm = grid_dm[grid];
738 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
739 Vec subX = globXArray[LAND_PACK_IDX(v_id, grid)], work = t_fhat[grid][tid];
740 ierr_t = PetscInfo(pack, "particlesToGrid %" PetscInt_FMT ".%" PetscInt_FMT ") for block %" PetscInt_FMT "\n", v_id, grid, LAND_PACK_IDX(v_id, grid));
741 ierr_t = particlesToGrid(dm, sw, tid, dim, wp_t[grid][tid], subX, globMpArray[LAND_PACK_IDX(v_id, grid)]);
742 if (ierr_t) ierr = ierr_t;
743 // u = M^_1 f_w
744 ierr_t = VecCopy(subX, work);
745 ierr_t = KSPSolve(t_ksp[grid][tid], work, subX);
746 if (ierr_t) ierr = ierr_t;
747 }
748 }
749 }
750 PetscCheck(!ierr, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Error in OMP loop. ierr = %d", (int)ierr);
751 /* Cleanup */
752 for (PetscInt tid = 0; tid < numthreads; tid++) {
753 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
754 if (glb_v_id < num_vertices) {
755 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
756 PetscCall(PetscFree4(xx_t[grid][tid], yy_t[grid][tid], wp_t[grid][tid], zz_t[grid][tid]));
757 }
758 } // active
759 } // threads
760 } // (fake) particle loop
761 // standard view of initial conditions
762 if (v_target >= global_vertex_id_0 && v_target < global_vertex_id_0 + ctx->batch_sz) {
763 PetscCall(DMSetOutputSequenceNumber(ctx->plex[g_target], 0, 0.0));
764 PetscCall(VecViewFromOptions(globXArray[LAND_PACK_IDX(v_target % ctx->batch_sz, g_target)], NULL, "-ex30_vec_view"));
765 if (ctx->num_grids > g_target + 1) {
766 PetscCall(DMSetOutputSequenceNumber(ctx->plex[g_target + 1], 0, 0.0));
767 PetscCall(VecViewFromOptions(globXArray[LAND_PACK_IDX(v_target % ctx->batch_sz, g_target + 1)], NULL, "-ex30_vec_view2"));
768 }
769 PetscCall(MatViewFromOptions(globMpArray[LAND_PACK_IDX(v_target % ctx->batch_sz, g_target)], NULL, "-ex30_mass_mat_view"));
770 PetscCall(DMViewFromOptions(globSwarmArray[LAND_PACK_IDX(v_target % ctx->batch_sz, g_target)], NULL, "-ex30_sw_view"));
771 PetscCall(DMSwarmViewXDMF(globSwarmArray[LAND_PACK_IDX(v_target % ctx->batch_sz, g_target)], "initial_swarm.xmf")); // writes a file by default!!!
772 }
773 // coarse graining moments_1a, bring f back from grid before advance
774 if (v_target >= global_vertex_id_0 && v_target < global_vertex_id_0 + ctx->batch_sz && printCtx->print_entropy) {
775 const PetscInt v_id = v_target % ctx->batch_sz;
776 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
777 PetscDataType dtype;
778 PetscReal *wp, *coords;
779 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
780 Vec work, subX = globXArray[LAND_PACK_IDX(v_id, grid)];
781 PetscInt bs, NN;
782 // C-G moments
783 PetscCall(VecDuplicate(subX, &work));
784 PetscCall(gridToParticles(grid_dm[grid], sw, subX, work, globMpArray[LAND_PACK_IDX(v_id, grid)], g_Mass[grid]));
785 PetscCall(VecDestroy(&work));
786 // moments
787 PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
788 PetscCall(DMSwarmGetLocalSize(sw, &NN));
789 PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wp));
790 for (PetscInt pp = 0; pp < NN; pp++) {
791 PetscReal v2 = 0, fact = (dim == 2) ? 2.0 * PETSC_PI * coords[pp * dim + 0] : 1, w = fact * wp[pp] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]];
792 for (PetscInt i = 0; i < dim; ++i) v2 += PetscSqr(coords[pp * dim + i]);
793 moments_1a[0] += w;
794 moments_1a[1] += w * ctx->v_0 * coords[pp * dim + 1]; // z-momentum
795 moments_1a[2] += w * 0.5 * ctx->v_0 * ctx->v_0 * v2;
796 }
797 PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
798 PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wp));
799 }
800 // entropy
801 const PetscReal N_inv = 1 / moments_1a[0];
802 PetscCall(PetscInfo(pack, "Entropy batch %" PetscInt_FMT " of %" PetscInt_FMT ", n = %g\n", v_target, num_vertices, (double)(1 / N_inv)));
803 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
804 PetscDataType dtype;
805 PetscReal *wp, *coords;
806 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
807 PetscInt bs, NN;
808 PetscCall(DMSwarmGetLocalSize(sw, &NN));
809 PetscCall(DMSwarmGetField(sw, "w_q", &bs, &dtype, (void **)&wp));
810 PetscCall(DMSwarmGetField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
811 for (PetscInt pp = 0; pp < NN; pp++) {
812 PetscReal fact = (dim == 2) ? 2.0 * PETSC_PI * coords[pp * dim + 0] : 1, w = fact * wp[pp] * ctx->n_0 * ctx->masses[ctx->species_offset[grid]], ww = w * N_inv;
813 if (w > PETSC_REAL_MIN) {
814 moments_1a[3] -= ww * PetscLogReal(ww);
815 PetscCheck(ww < 1 - PETSC_MACHINE_EPSILON, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "ww (%g) > 1", (double)ww);
816 } else moments_1a[4] -= w;
817 }
818 PetscCall(DMSwarmRestoreField(sw, "w_q", &bs, &dtype, (void **)&wp));
819 PetscCall(DMSwarmRestoreField(sw, "DMSwarmPIC_coor", &bs, &dtype, (void **)&coords));
820 }
821 }
822 // restore vector
823 PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, globXArray));
824 // view initial grid
825 if (v_target >= global_vertex_id_0 && v_target < global_vertex_id_0 + ctx->batch_sz) PetscCall(DMPlexLandauPrintNorms(X, 0));
826 // advance
827 PetscCall(TSSetSolution(ts, X));
828 PetscCall(PetscInfo(pack, "Advance vertex %" PetscInt_FMT " to %" PetscInt_FMT "\n", global_vertex_id_0, global_vertex_id_0 + ctx->batch_sz));
829 PetscCall(TSSetPostStep(ts, PostStep));
830 PetscCall(PostStep(ts));
831 PetscCall(TSSolve(ts, X));
832 // view
833 PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, globXArray));
834 if (v_target >= global_vertex_id_0 && v_target < global_vertex_id_0 + ctx->batch_sz) {
835 /* Visualize original particle field */
836 DM sw = globSwarmArray[LAND_PACK_IDX(v_target % ctx->batch_sz, g_target)];
837 Vec f;
838 PetscCall(DMSetOutputSequenceNumber(sw, 0, 0.0));
839 PetscCall(DMViewFromOptions(grid_dm[g_target], NULL, "-weights_dm_view"));
840 PetscCall(DMViewFromOptions(sw, NULL, "-weights_sw_view"));
841 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
842 PetscCall(PetscObjectSetName((PetscObject)f, "weights"));
843 PetscCall(VecViewFromOptions(f, NULL, "-weights_vec_view"));
844 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
845 //
846 PetscCall(DMPlexLandauPrintNorms(X, 1));
847 }
848 if (!use_uniform_particle_grid) { // resample to uniform grid
849 for (PetscInt v_id_0 = 0; v_id_0 < ctx->batch_sz; v_id_0 += numthreads) {
850 PetscReal *xx_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS], *yy_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS], *zz_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS], *wp_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS];
851 PetscInt Np_t[LANDAU_MAX_GRIDS][EX30_MAX_NUM_THRDS];
852 for (PetscInt tid = 0; tid < numthreads; tid++) {
853 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
854 if (glb_v_id < num_vertices) {
855 // create uniform grid w/o weights & smaller
856 PetscInt Npp0 = (a_Np + (glb_v_id % (a_Np / 10 + 1))) / 2, Nv; // 1/2 of uniform particle grid size
857 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
858 // for (PetscInt sp = ctx->species_offset[grid], i0 = 0; sp < ctx->species_offset[grid + 1]; sp++, i0++)
859 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];
860 PetscInt Npi = Npp0, Npj = 2 * Npp0, Npk = 1, NN;
861 // delete old particles and particle mass matrix
862 PetscCall(DMDestroy(&globSwarmArray[LAND_PACK_IDX(v_id, grid)]));
863 PetscCall(MatDestroy(&globMpArray[LAND_PACK_IDX(v_id, grid)]));
864 // create fake particles in batches with threads
865 PetscCall(MatGetLocalSize(g_Mass[grid], &Nv, NULL));
866 if (dim == 2) lo[0] = 0;
867 else Npi = Npj = Npk = Npp0;
868 NN = Npi * Npj * Npk + (dim == 2 ? 3 : 6); // make a regular grid of particles Npp x Npp
869 while (Npi * Npj * Npk < Nv) { // make stable - no LS
870 Npi++;
871 Npj++;
872 Npk++;
873 NN = Npi * Npj * Npk + (dim == 2 ? 3 : 6);
874 }
875 Np_t[grid][tid] = NN;
876 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]));
877 hp[0] = (hi[0] - lo[0]) / Npi;
878 hp[1] = (hi[1] - lo[1]) / Npj;
879 hp[2] = (hi[2] - lo[2]) / Npk;
880 if (dim == 2) hp[2] = 1;
881 PetscCall(PetscInfo(pack, "Resampling %d particles, %d vertices\n", (int)NN, (int)Nv)); // temp
882 for (PetscInt pj = 0, pp = 0; pj < Npj; pj++) {
883 for (PetscInt pk = 0; pk < Npk; pk++) {
884 for (PetscInt pi = 0; pi < Npi; pi++, pp++) {
885 wp_t[grid][tid][pp] = 0;
886 xx_t[grid][tid][pp] = lo[0] + hp[0] / 2.0 + pi * hp[0];
887 yy_t[grid][tid][pp] = lo[1] + hp[1] / 2.0 + pj * hp[1];
888 if (dim == 3) zz_t[grid][tid][pp] = lo[2] + hp[2] / 2.0 + pk * hp[2];
889 }
890 }
891 }
892 if (dim == 2) { // fix bounding box
893 PetscInt pp = NN - 3;
894 wp_t[grid][tid][pp] = 0;
895 xx_t[grid][tid][pp] = 1.e-7;
896 yy_t[grid][tid][pp++] = hi[1] - 5.e-7;
897 wp_t[grid][tid][pp] = 0;
898 xx_t[grid][tid][pp] = hi[0] - 5.e-7;
899 yy_t[grid][tid][pp++] = 0;
900 wp_t[grid][tid][pp] = 0;
901 xx_t[grid][tid][pp] = 1.e-7;
902 yy_t[grid][tid][pp++] = lo[1] + 5.e-7;
903 } else {
904 const PetscInt p0 = NN - 6;
905 for (PetscInt pj = 0; pj < 6; pj++) xx_t[grid][tid][p0 + pj] = yy_t[grid][tid][p0 + pj] = zz_t[grid][tid][p0 + pj] = wp_t[grid][tid][p0 + pj] = 0;
906 xx_t[grid][tid][p0 + 0] = lo[0];
907 xx_t[grid][tid][p0 + 1] = hi[0];
908 yy_t[grid][tid][p0 + 2] = lo[1];
909 yy_t[grid][tid][p0 + 3] = hi[1];
910 zz_t[grid][tid][p0 + 4] = lo[2];
911 zz_t[grid][tid][p0 + 5] = hi[2];
912 }
913 }
914 } // active
915 } // threads
916 /* Create particle swarm */
917 for (PetscInt tid = 0; tid < numthreads; tid++) {
918 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
919 if (glb_v_id < num_vertices) { // the ragged edge of the last batch
920 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
921 // for (PetscInt sp = ctx->species_offset[grid], i0 = 0; sp < ctx->species_offset[grid + 1]; sp++, i0++) -- loop over species for Nf > 1 -- TODO
922 PetscErrorCode ierr_t;
923 PetscSection section;
924 PetscInt Nf;
925 DM dm = grid_dm[grid];
926 ierr_t = DMGetLocalSection(dm, §ion);
927 ierr_t = PetscSectionGetNumFields(section, &Nf);
928 if (Nf != 1) ierr_t = (PetscErrorCode)9999;
929 else {
930 ierr_t = DMViewFromOptions(dm, NULL, "-dm_view");
931 ierr_t = PetscInfo(pack, "call createSwarm [%" PetscInt_FMT ".%" PetscInt_FMT "] local block index %" PetscInt_FMT "\n", v_id, grid, LAND_PACK_IDX(v_id, grid));
932 ierr_t = createSwarm(dm, dim, &globSwarmArray[LAND_PACK_IDX(v_id, grid)]);
933 }
934 if (ierr_t) ierr = ierr_t;
935 }
936 } // active
937 } // threads
938 PetscCheck(ierr != 9999, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Only support one species per grid");
939 PetscCheck(!ierr, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Error in OMP loop. ierr = %d", (int)ierr);
940 // make globMpArray
941 PetscPragmaOMP(parallel for)
942 for (PetscInt tid = 0; tid < numthreads; tid++) {
943 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
944 if (glb_v_id < num_vertices) {
945 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
946 // for (PetscInt sp = ctx->species_offset[grid], i0 = 0; sp < ctx->species_offset[grid + 1]; sp++, i0++) -- loop over species for Nf > 1 -- TODO
947 PetscErrorCode ierr_t;
948 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
949 ierr_t = PetscInfo(pack, "makeSwarm %" PetscInt_FMT ".%" PetscInt_FMT ") for block %" PetscInt_FMT "\n", v_id, grid, LAND_PACK_IDX(v_id, grid));
950 ierr_t = makeSwarm(sw, dim, Np_t[grid][tid], xx_t[grid][tid], yy_t[grid][tid], zz_t[grid][tid]);
951 if (ierr_t) ierr = ierr_t;
952 }
953 } // active
954 } // threads
955 // create particle mass matrices
956 //PetscPragmaOMP(parallel for)
957 for (PetscInt tid = 0; tid < numthreads; tid++) {
958 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
959 if (glb_v_id < num_vertices) {
960 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
961 PetscErrorCode ierr_t;
962 DM dm = grid_dm[grid];
963 DM sw = globSwarmArray[LAND_PACK_IDX(v_id, grid)];
964 ierr_t = PetscInfo(pack, "createMp %" PetscInt_FMT ".%" PetscInt_FMT ") for block %" PetscInt_FMT "\n", v_id, grid, LAND_PACK_IDX(v_id, grid));
965 ierr_t = createMp(dm, sw, &globMpArray[LAND_PACK_IDX(v_id, grid)]);
966 if (ierr_t) ierr = ierr_t;
967 }
968 } // active
969 } // threads
970 PetscCheck(!ierr, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Error in OMP loop. ierr = %d", (int)ierr);
971 /* Cleanup */
972 for (PetscInt tid = 0; tid < numthreads; tid++) {
973 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
974 if (glb_v_id < num_vertices) {
975 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
976 PetscCall(PetscFree4(xx_t[grid][tid], yy_t[grid][tid], wp_t[grid][tid], zz_t[grid][tid]));
977 }
978 } // active
979 } // threads
980 } // batch
981 // view
982 if (v_target >= global_vertex_id_0 && v_target < global_vertex_id_0 + ctx->batch_sz) {
983 /* Visualize particle field */
984 DM sw = globSwarmArray[LAND_PACK_IDX(v_target % ctx->batch_sz, g_target)];
985 Vec f;
986 PetscCall(DMSetOutputSequenceNumber(sw, 0, 0.0));
987 PetscCall(DMViewFromOptions(sw, NULL, "-resampled_weights_sw_view"));
988 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f));
989 PetscCall(PetscObjectSetName((PetscObject)f, "resampled_weights"));
990 PetscCall(VecViewFromOptions(f, NULL, "-resampled_weights_vec_view"));
991 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f));
992 PetscCall(DMSwarmViewXDMF(sw, "resampled.xmf"));
993 }
994 } // !uniform
995 // particles to grid, compute moments and entropy, for target vertex only
996 if (v_target >= global_vertex_id_0 && v_target < global_vertex_id_0 + ctx->batch_sz && printCtx->print_entropy) {
997 PetscReal energy_error_rel;
998 PetscCall(gridToParticles_private(grid_dm, globSwarmArray, dim, v_target, numthreads, num_vertices, global_vertex_id_0, globMpArray, g_Mass, t_fhat, moments_1b, globXArray, ctx));
999 energy_error_rel = PetscAbsReal(moments_1b[2] - moments_0[2]) / moments_0[2];
1000 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Particle Moments:\t number density momentum (par) energy entropy negative weights : # OMP threads %g\n", (double)numthreads));
1001 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\tInitial: %18.12e %19.12e %18.12e %18.12e %g %%\n", (double)moments_0[0], (double)moments_0[1], (double)moments_0[2], (double)moments_0[3], 100 * (double)(moments_0[4] / moments_0[0])));
1002 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\tCoarse-graining: %18.12e %19.12e %18.12e %18.12e %g %%\n", (double)moments_1a[0], (double)moments_1a[1], (double)moments_1a[2], (double)moments_1a[3], 100 * (double)(moments_1a[4] / moments_0[0])));
1003 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\tLandau: %18.12e %19.12e %18.12e %18.12e %g %%\n", (double)moments_1b[0], (double)moments_1b[1], (double)moments_1b[2], (double)moments_1b[3], 100 * (double)(moments_1b[4] / moments_0[0])));
1004 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Coarse-graining entropy generation = %e ; Landau entropy generation = %e\n", (double)(moments_1a[3] - moments_0[3]), (double)(moments_1b[3] - moments_0[3])));
1005 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "(relative) energy conservation: Coarse-graining = %e, Landau = %e (%g %d)\n", (double)(moments_1a[2] - moments_0[2]) / (double)moments_0[2], (double)energy_error_rel, (double)PetscLog10Real(energy_error_rel), (int)(PetscLog10Real(energy_error_rel) + .5)));
1006 }
1007 // restore vector
1008 PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, globXArray));
1009 // cleanup
1010 for (PetscInt v_id_0 = 0; v_id_0 < ctx->batch_sz; v_id_0 += numthreads) {
1011 for (PetscInt tid = 0; tid < numthreads; tid++) {
1012 const PetscInt v_id = v_id_0 + tid, glb_v_id = global_vertex_id_0 + v_id;
1013 if (glb_v_id < num_vertices) {
1014 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
1015 PetscCall(DMDestroy(&globSwarmArray[LAND_PACK_IDX(v_id, grid)]));
1016 PetscCall(MatDestroy(&globMpArray[LAND_PACK_IDX(v_id, grid)]));
1017 }
1018 }
1019 }
1020 }
1021 } // user batch, not used
1022 /* Cleanup */
1023 PetscCall(PetscFree(globXArray));
1024 PetscCall(PetscFree(globSwarmArray));
1025 PetscCall(PetscFree(globMpArray));
1026 PetscCall(PetscFree(printCtx));
1027 // clean up mass matrices
1028 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // add same particels for all grids
1029 PetscCall(MatDestroy(&g_Mass[grid]));
1030 for (PetscInt tid = 0; tid < numthreads; tid++) {
1031 PetscCall(VecDestroy(&t_fhat[grid][tid]));
1032 PetscCall(KSPDestroy(&t_ksp[grid][tid]));
1033 }
1034 }
1035 PetscFunctionReturn(PETSC_SUCCESS);
1036 }
1037
main(int argc,char ** argv)1038 int main(int argc, char **argv)
1039 {
1040 DM pack;
1041 Vec X;
1042 PetscInt dim = 2, num_vertices = 1, Np = 10, v_target = 0, g_target = 0;
1043 TS ts;
1044 Mat J;
1045 LandauCtx *ctx;
1046 PetscReal shift = 0;
1047 PetscBool use_uniform_particle_grid = PETSC_TRUE;
1048
1049 PetscFunctionBeginUser;
1050 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
1051 // process args
1052 PetscOptionsBegin(PETSC_COMM_SELF, "", "Collision Options", "DMPLEX");
1053 PetscCall(PetscOptionsInt("-dim", "Velocity space dimension", "ex30.c", dim, &dim, NULL));
1054 PetscCall(PetscOptionsInt("-number_spatial_vertices", "Number of user spatial vertices to be batched for Landau", "ex30.c", num_vertices, &num_vertices, NULL));
1055 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));
1056 PetscCall(PetscOptionsBool("-use_uniform_particle_grid", "Use uniform particle grid", "ex30.c", use_uniform_particle_grid, &use_uniform_particle_grid, NULL));
1057 PetscCall(PetscOptionsInt("-vertex_view_target", "Global vertex for diagnostics", "ex30.c", v_target, &v_target, NULL));
1058 PetscCall(PetscOptionsReal("-e_shift", "Bi-Maxwellian shift", "ex30.c", shift, &shift, NULL));
1059 PetscCheck(v_target < num_vertices, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Batch to view %" PetscInt_FMT " should be < number of vertices %" PetscInt_FMT, v_target, num_vertices);
1060 PetscCall(PetscOptionsInt("-grid_view_target", "Grid to view with diagnostics", "ex30.c", g_target, &g_target, NULL));
1061 PetscOptionsEnd();
1062 /* Create a mesh */
1063 PetscCall(DMPlexLandauCreateVelocitySpace(PETSC_COMM_SELF, dim, "", &X, &J, &pack));
1064 PetscCall(DMGetApplicationContext(pack, &ctx));
1065 PetscCall(DMSetUp(pack));
1066 PetscCall(DMSetOutputSequenceNumber(pack, 0, 0.0));
1067 PetscCheck(g_target < ctx->num_grids, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Grid to view %" PetscInt_FMT " should be < number of grids %" PetscInt_FMT, g_target, ctx->num_grids);
1068 PetscCheck(ctx->batch_view_idx == v_target % ctx->batch_sz, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Global view index %" PetscInt_FMT " mode batch size %" PetscInt_FMT " != ctx->batch_view_idx %" PetscInt_FMT, v_target, ctx->batch_sz, ctx->batch_view_idx);
1069 // PetscCheck(!use_uniform_particle_grid || !ctx->sphere, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Can not use -use_uniform_particle_grid and -dm_landau_sphere");
1070 /* Create timestepping solver context */
1071 PetscCall(TSCreate(PETSC_COMM_SELF, &ts));
1072 PetscCall(TSSetDM(ts, pack));
1073 PetscCall(TSSetIFunction(ts, NULL, DMPlexLandauIFunction, NULL));
1074 PetscCall(TSSetIJacobian(ts, J, J, DMPlexLandauIJacobian, NULL));
1075 PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
1076 PetscCall(TSSetFromOptions(ts));
1077 PetscCall(PetscObjectSetName((PetscObject)X, "X"));
1078 // do particle advance
1079 PetscCall(go(ts, X, num_vertices, Np, dim, v_target, g_target, shift, use_uniform_particle_grid));
1080 PetscCall(MatZeroEntries(J)); // need to zero out so as to not reuse it in Landau's logic
1081 /* clean up */
1082 PetscCall(DMPlexLandauDestroyVelocitySpace(&pack));
1083 PetscCall(TSDestroy(&ts));
1084 PetscCall(VecDestroy(&X));
1085 PetscCall(PetscFinalize());
1086 return 0;
1087 }
1088
1089 /*TEST
1090
1091 build:
1092 requires: !complex
1093
1094 testset:
1095 requires: double defined(PETSC_USE_DMLANDAU_2D)
1096 output_file: output/ex30_0.out
1097 args: -dim 2 -petscspace_degree 3 -dm_landau_num_species_grid 1,1,1 -dm_refine 1 -number_particles_per_dimension 20 \
1098 -dm_landau_batch_size 4 -number_spatial_vertices 6 -vertex_view_target 5 -grid_view_target 1 -dm_landau_batch_view_idx 1 \
1099 -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 \
1100 -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 -ftop_ksp_error_if_not_converged \
1101 -ksp_type gmres -ksp_error_if_not_converged -dm_landau_verbose 4 -print_entropy \
1102 -ptof_ksp_type cg -ptof_pc_type jacobi -ptof_ksp_rtol 1e-12 -ptof_ksp_error_if_not_converged\
1103 -snes_converged_reason -snes_monitor -snes_rtol 1e-12 -snes_stol 1e-12 \
1104 -ts_time_step 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
1105 test:
1106 suffix: cpu
1107 args: -dm_landau_device_type cpu -pc_type jacobi
1108 test:
1109 suffix: kokkos
1110 # failed on Sunspot@ALCF with sycl
1111 requires: kokkos_kernels !openmp !sycl
1112 args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -pc_type bjkokkos -pc_bjkokkos_ksp_type tfqmr -pc_bjkokkos_pc_type jacobi
1113
1114 testset:
1115 requires: double !defined(PETSC_USE_DMLANDAU_2D)
1116 output_file: output/ex30_3d.out
1117 args: -dim 3 -petscspace_degree 2 -dm_landau_num_species_grid 1,1 -dm_refine 0 -number_particles_per_dimension 10 -dm_plex_hash_location \
1118 -dm_landau_batch_size 1 -number_spatial_vertices 1 -vertex_view_target 0 -grid_view_target 0 -dm_landau_batch_view_idx 0 \
1119 -dm_landau_n 1.000018,1 -dm_landau_thermal_temps 2,1 -dm_landau_ion_masses 2 -dm_landau_ion_charges 1 \
1120 -ftop_ksp_type cg -ftop_pc_type jacobi -ftop_ksp_rtol 1e-12 -ftop_ksp_error_if_not_converged -ksp_type preonly -pc_type lu -ksp_error_if_not_converged \
1121 -ptof_ksp_type cg -ptof_pc_type jacobi -ptof_ksp_rtol 1e-12 -ptof_ksp_error_if_not_converged \
1122 -snes_converged_reason -snes_monitor -snes_rtol 1e-12 -snes_stol 1e-12 \
1123 -ts_time_step 0.1 -ts_exact_final_time stepover -ts_max_snes_failures -1 -ts_max_steps 1 -ts_monitor -ts_type beuler -print_entropy
1124 test:
1125 suffix: cpu_3d
1126 args: -dm_landau_device_type cpu
1127 test:
1128 suffix: kokkos_3d
1129 requires: kokkos_kernels !openmp
1130 args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -pc_type bjkokkos -pc_bjkokkos_ksp_type tfqmr -pc_bjkokkos_pc_type jacobi
1131
1132 test:
1133 suffix: conserve
1134 requires: !complex double defined(PETSC_USE_DMLANDAU_2D) !cuda
1135 args: -dm_landau_batch_size 4 -dm_refine 0 -dm_landau_num_species_grid 1 -dm_landau_thermal_temps 1 -petscspace_degree 3 -snes_converged_reason -ts_type beuler -ts_time_step .1 \
1136 -ts_max_steps 1 -ksp_type preonly -ksp_error_if_not_converged -snes_rtol 1e-14 -snes_stol 1e-14 -dm_landau_device_type cpu -number_particles_per_dimension 20 \
1137 -ptof_ksp_type cg -ptof_pc_type jacobi -ptof_ksp_rtol 1e-14 -ptof_ksp_error_if_not_converged -pc_type lu -dm_landau_simplex 1 -use_uniform_particle_grid false -dm_landau_sphere -print_entropy -number_particles_per_dimension 50 -ftop_ksp_type cg -ftop_pc_type jacobi -ftop_ksp_rtol 1e-14
1138
1139 TEST*/
1140