Actual source code: ex2.c

  1: static char help[] = "Runaway electron model with Landau collision operator\n\n";

  3: #include <petscdmplex.h>
  4: #include <petsclandau.h>
  5: #include <petscts.h>
  6: #include <petscds.h>
  7: #include <petscdmcomposite.h>
  8: #include "petsc/private/petscimpl.h"

 10: #if defined(PETSC_HAVE_CUDA_NVTX)
 11:   #include <nvToolsExt.h>
 12: #endif

 14: /* data for runaway electron model */
 15: typedef struct REctx_struct {
 16:   PetscErrorCode (*test)(TS, Vec, PetscInt, PetscReal, PetscBool, LandauCtx *, struct REctx_struct *);
 17:   PetscErrorCode (*impuritySrcRate)(PetscReal, PetscReal *, LandauCtx *);
 18:   PetscErrorCode (*E)(Vec, Vec, PetscInt, PetscReal, LandauCtx *, PetscReal *);
 19:   PetscReal T_cold;        /* temperature of newly ionized electrons and impurity ions */
 20:   PetscReal ion_potential; /* ionization potential of impurity */
 21:   PetscReal Ne_ion;        /* effective number of electrons shed in ioization of impurity */
 22:   PetscReal Ez_initial;
 23:   PetscReal L; /* inductance */
 24:   Vec       X_0;
 25:   PetscInt  imp_idx; /* index for impurity ionizing sink */
 26:   PetscReal pulse_start;
 27:   PetscReal pulse_width;
 28:   PetscReal pulse_rate;
 29:   PetscReal current_rate;
 30:   PetscInt  plotIdx;
 31:   PetscInt  plotStep;
 32:   PetscInt  idx; /* cache */
 33:   PetscReal j;   /* cache */
 34:   PetscReal plotDt;
 35:   PetscBool plotting;
 36:   PetscBool use_spitzer_eta;
 37:   PetscInt  print_period;
 38:   PetscInt  grid_view_idx;
 39: } REctx;

 41: static const PetscReal kev_joul = 6.241506479963235e+15; /* 1/1000e */

 43: #define RE_CUT 3.
 44: /* < v, u_re * v * q > */
 45: static void f0_j_re(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0)
 46: {
 47:   PetscReal n_e = PetscRealPart(u[0]);
 48:   if (dim == 2) {
 49:     if (x[1] > RE_CUT || x[1] < -RE_CUT) {                    /* simply a cutoff for REs. v_|| > 3 v(T_e) */
 50:       *f0 = n_e * 2. * PETSC_PI * x[0] * x[1] * constants[0]; /* n * r * v_|| * q */
 51:     } else {
 52:       *f0 = 0;
 53:     }
 54:   } else {
 55:     if (x[2] > RE_CUT || x[2] < -RE_CUT) { /* simply a cutoff for REs. v_|| > 3 v(T_e) */
 56:       *f0 = n_e * x[2] * constants[0];
 57:     } else {
 58:       *f0 = 0;
 59:     }
 60:   }
 61: }

 63: /* sum < v, u*v*q > */
 64: static void f0_jz_sum(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar q[], PetscScalar *f0)
 65: {
 66:   PetscInt ii;
 67:   f0[0] = 0;
 68:   if (dim == 2) {
 69:     for (ii = 0; ii < Nf; ii++) f0[0] += u[ii] * 2. * PETSC_PI * x[0] * x[1] * q[ii]; /* n * r * v_|| * q * v_0 */
 70:   } else {
 71:     for (ii = 0; ii < Nf; ii++) f0[0] += u[ii] * x[2] * q[ii]; /* n * v_|| * q  * v_0 */
 72:   }
 73: }

 75: /* < v, n_e > */
 76: static void f0_n(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0)
 77: {
 78:   PetscInt ii = (PetscInt)PetscRealPart(constants[0]);
 79:   if (dim == 2) f0[0] = 2. * PETSC_PI * x[0] * u[ii];
 80:   else f0[0] = u[ii];
 81: }

 83: /* < v, n_e v_|| > */
 84: static void f0_vz(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0)
 85: {
 86:   PetscInt ii = (PetscInt)PetscRealPart(constants[0]);
 87:   if (dim == 2) f0[0] = u[ii] * 2. * PETSC_PI * x[0] * x[1]; /* n r v_|| */
 88:   else f0[0] = u[ii] * x[2];                                 /* n v_|| */
 89: }

 91: /* < v, n_e (v-shift) > */
 92: static void f0_ve_shift(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0)
 93: {
 94:   PetscReal vz = numConstants > 0 ? PetscRealPart(constants[0]) : 0;
 95:   if (dim == 2) *f0 = u[0] * 2. * PETSC_PI * x[0] * PetscSqrtReal(x[0] * x[0] + (x[1] - vz) * (x[1] - vz)); /* n r v */
 96:   else {
 97:     *f0 = u[0] * PetscSqrtReal(x[0] * x[0] + x[1] * x[1] + (x[2] - vz) * (x[2] - vz)); /* n v */
 98:   }
 99: }

101: /* CalculateE - Calculate the electric field  */
102: /*  T        -- Electron temperature  */
103: /*  n        -- Electron density  */
104: /*  lnLambda --   */
105: /*  eps0     --  */
106: /*  E        -- output E, input \hat E */
107: static PetscReal CalculateE(PetscReal Tev, PetscReal n, PetscReal lnLambda, PetscReal eps0, PetscReal *E)
108: {
109:   PetscReal c, e, m;

111:   c = 299792458.0;
112:   e = 1.602176e-19;
113:   m = 9.10938e-31;
114:   if (1) {
115:     double Ec, Ehat = *E, betath = PetscSqrtReal(2 * Tev * e / (m * c * c)), j0 = Ehat * 7 / (PetscSqrtReal(2) * 2) * PetscPowReal(betath, 3) * n * e * c;
116:     Ec = n * lnLambda * PetscPowReal(e, 3) / (4 * PETSC_PI * PetscPowReal(eps0, 2) * m * c * c);
117:     *E = Ec;
118:     PetscPrintf(PETSC_COMM_WORLD, "CalculateE j0=%g Ec = %g\n", j0, Ec);
119:   } else {
120:     PetscReal Ed, vth;
121:     vth = PetscSqrtReal(8 * Tev * e / (m * PETSC_PI));
122:     Ed  = n * lnLambda * PetscPowReal(e, 3) / (4 * PETSC_PI * PetscPowReal(eps0, 2) * m * vth * vth);
123:     *E  = Ed;
124:   }
125:   return 0;
126: }

128: static PetscReal Spitzer(PetscReal m_e, PetscReal e, PetscReal Z, PetscReal epsilon0, PetscReal lnLam, PetscReal kTe_joules)
129: {
130:   PetscReal Fz = (1 + 1.198 * Z + 0.222 * Z * Z) / (1 + 2.966 * Z + 0.753 * Z * Z), eta;
131:   eta          = Fz * 4. / 3. * PetscSqrtReal(2. * PETSC_PI) * Z * PetscSqrtReal(m_e) * PetscSqr(e) * lnLam * PetscPowReal(4 * PETSC_PI * epsilon0, -2.) * PetscPowReal(kTe_joules, -1.5);
132:   return eta;
133: }

135: /*  */
136: static PetscErrorCode testNone(TS ts, Vec X, PetscInt stepi, PetscReal time, PetscBool islast, LandauCtx *ctx, REctx *rectx)
137: {
139:   return 0;
140: }

142: /*  */
143: static PetscErrorCode testSpitzer(TS ts, Vec X, PetscInt stepi, PetscReal time, PetscBool islast, LandauCtx *ctx, REctx *rectx)
144: {
145:   PetscInt          ii, nDMs;
146:   PetscDS           prob;
147:   static PetscReal  old_ratio = 1e10;
148:   TSConvergedReason reason;
149:   PetscReal         J, J_re, spit_eta, Te_kev = 0, E, ratio, Z, n_e, v, v2;
150:   PetscScalar       user[2] = {0., ctx->charges[0]}, q[LANDAU_MAX_SPECIES], tt[LANDAU_MAX_SPECIES], vz;
151:   PetscReal         dt;
152:   DM                pack, plexe = ctx->plex[0], plexi = (ctx->num_grids == 1) ? NULL : ctx->plex[1];
153:   Vec              *XsubArray;

157:   VecGetDM(X, &pack);
159:   DMCompositeGetNumberDM(pack, &nDMs);
161:   PetscMalloc(sizeof(*XsubArray) * nDMs, &XsubArray);
162:   DMCompositeGetAccessArray(pack, X, nDMs, NULL, XsubArray); // read only
163:   TSGetTimeStep(ts, &dt);
164:   /* get current for each grid */
165:   for (ii = 0; ii < ctx->num_species; ii++) q[ii] = ctx->charges[ii];
166:   DMGetDS(plexe, &prob);
167:   PetscDSSetConstants(prob, 2, &q[0]);
168:   PetscDSSetObjective(prob, 0, &f0_jz_sum);
169:   DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL);
170:   J = -ctx->n_0 * ctx->v_0 * PetscRealPart(tt[0]);
171:   if (plexi) { // add first (only) ion
172:     DMGetDS(plexi, &prob);
173:     PetscDSSetConstants(prob, 1, &q[1]);
174:     PetscDSSetObjective(prob, 0, &f0_jz_sum);
175:     DMPlexComputeIntegralFEM(plexi, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 1)], tt, NULL);
176:     J += -ctx->n_0 * ctx->v_0 * PetscRealPart(tt[0]);
177:   }
178:   /* get N_e */
179:   DMGetDS(plexe, &prob);
180:   PetscDSSetConstants(prob, 1, user);
181:   PetscDSSetObjective(prob, 0, &f0_n);
182:   DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL);
183:   n_e = PetscRealPart(tt[0]) * ctx->n_0;
184:   /* Z */
185:   Z = -ctx->charges[1] / ctx->charges[0];
186:   /* remove drift */
187:   if (0) {
188:     user[0] = 0; // electrons
189:     DMGetDS(plexe, &prob);
190:     PetscDSSetConstants(prob, 1, user);
191:     PetscDSSetObjective(prob, 0, &f0_vz);
192:     DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL);
193:     vz = ctx->n_0 * PetscRealPart(tt[0]) / n_e; /* non-dimensional */
194:   } else vz = 0;
195:   /* thermal velocity */
196:   DMGetDS(plexe, &prob);
197:   PetscDSSetConstants(prob, 1, &vz);
198:   PetscDSSetObjective(prob, 0, &f0_ve_shift);
199:   DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL);
200:   v        = ctx->n_0 * ctx->v_0 * PetscRealPart(tt[0]) / n_e;                                           /* remove number density to get velocity */
201:   v2       = PetscSqr(v);                                                                                /* use real space: m^2 / s^2 */
202:   Te_kev   = (v2 * ctx->masses[0] * PETSC_PI / 8) * kev_joul;                                            /* temperature in kev */
203:   spit_eta = Spitzer(ctx->masses[0], -ctx->charges[0], Z, ctx->epsilon0, ctx->lnLam, Te_kev / kev_joul); /* kev --> J (kT) */
204:   if (0) {
205:     DMGetDS(plexe, &prob);
206:     PetscDSSetConstants(prob, 1, q);
207:     PetscDSSetObjective(prob, 0, &f0_j_re);
208:     DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL);
209:   } else tt[0] = 0;
210:   J_re = -ctx->n_0 * ctx->v_0 * PetscRealPart(tt[0]);
211:   DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, XsubArray); // read only
212:   PetscFree(XsubArray);

214:   if (rectx->use_spitzer_eta) {
215:     E = ctx->Ez = spit_eta * (rectx->j - J_re);
216:   } else {
217:     E        = ctx->Ez; /* keep real E */
218:     rectx->j = J;       /* cache */
219:   }

221:   ratio = E / J / spit_eta;
222:   if (stepi > 10 && !rectx->use_spitzer_eta && ((old_ratio - ratio < 1.e-6))) {
223:     rectx->pulse_start     = time + 0.98 * dt;
224:     rectx->use_spitzer_eta = PETSC_TRUE;
225:   }
226:   TSGetConvergedReason(ts, &reason);
227:   TSGetConvergedReason(ts, &reason);
228:   if ((rectx->plotting) || stepi == 0 || reason || rectx->pulse_start == time + 0.98 * dt) {
229:     PetscCall(PetscPrintf(ctx->comm, "testSpitzer: %4" PetscInt_FMT ") time=%11.4e n_e= %10.3e E= %10.3e J= %10.3e J_re= %10.3e %.3g%% Te_kev= %10.3e Z_eff=%g E/J to eta ratio= %g (diff=%g) %s %s spit_eta=%g\n", stepi, (double)time,
230:                           (double)(n_e / ctx->n_0), (double)ctx->Ez, (double)J, (double)J_re, (double)(100 * J_re / J), (double)Te_kev, (double)Z, (double)ratio, (double)(old_ratio - ratio), rectx->use_spitzer_eta ? "using Spitzer eta*J E" : "constant E", rectx->pulse_start != time + 0.98 * dt ? "normal" : "transition", (double)spit_eta));
232:   }
233:   old_ratio = ratio;
234:   return 0;
235: }

237: static const double ppp = 2;
238: static void f0_0_diff_lp(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0)
239: {
240:   LandauCtx      *ctx   = (LandauCtx *)constants;
241:   REctx          *rectx = (REctx *)ctx->data;
242:   PetscInt        ii    = rectx->idx, i;
243:   const PetscReal kT_m  = ctx->k * ctx->thermal_temps[ii] / ctx->masses[ii]; /* kT/m */
244:   const PetscReal n     = ctx->n[ii];
245:   PetscReal       diff, f_maxwell, v2 = 0, theta = 2 * kT_m / (ctx->v_0 * ctx->v_0); /* theta = 2kT/mc^2 */
246:   for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
247:   f_maxwell = n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
248:   diff      = 2. * PETSC_PI * x[0] * (PetscRealPart(u[ii]) - f_maxwell);
249:   f0[0]     = PetscPowReal(diff, ppp);
250: }
251: static void f0_0_maxwellian_lp(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0)
252: {
253:   LandauCtx      *ctx   = (LandauCtx *)constants;
254:   REctx          *rectx = (REctx *)ctx->data;
255:   PetscInt        ii    = rectx->idx, i;
256:   const PetscReal kT_m  = ctx->k * ctx->thermal_temps[ii] / ctx->masses[ii]; /* kT/m */
257:   const PetscReal n     = ctx->n[ii];
258:   PetscReal       f_maxwell, v2 = 0, theta = 2 * kT_m / (ctx->v_0 * ctx->v_0); /* theta = 2kT/mc^2 */
259:   for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
260:   f_maxwell = 2. * PETSC_PI * x[0] * n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
261:   f0[0]     = PetscPowReal(f_maxwell, ppp);
262: }

264: /*  */
265: static PetscErrorCode testStable(TS ts, Vec X, PetscInt stepi, PetscReal time, PetscBool islast, LandauCtx *ctx, REctx *rectx)
266: {
267:   PetscDS     prob;
268:   Vec         X2;
269:   PetscReal   ediff, idiff = 0, lpm0, lpm1 = 1;
270:   PetscScalar tt[LANDAU_MAX_SPECIES];
271:   DM          dm, plex = ctx->plex[0];

274:   VecGetDM(X, &dm);
275:   DMGetDS(plex, &prob);
276:   VecDuplicate(X, &X2);
277:   VecCopy(X, X2);
278:   if (!rectx->X_0) {
279:     VecDuplicate(X, &rectx->X_0);
280:     VecCopy(X, rectx->X_0);
281:   }
282:   VecAXPY(X, -1.0, rectx->X_0);
283:   PetscDSSetConstants(prob, sizeof(LandauCtx) / sizeof(PetscScalar), (PetscScalar *)ctx);
284:   rectx->idx = 0;
285:   PetscDSSetObjective(prob, 0, &f0_0_diff_lp);
286:   DMPlexComputeIntegralFEM(plex, X2, tt, NULL);
287:   ediff = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
288:   PetscDSSetObjective(prob, 0, &f0_0_maxwellian_lp);
289:   DMPlexComputeIntegralFEM(plex, X2, tt, NULL);
290:   lpm0 = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
291:   if (ctx->num_species > 1) {
292:     rectx->idx = 1;
293:     PetscDSSetObjective(prob, 0, &f0_0_diff_lp);
294:     DMPlexComputeIntegralFEM(plex, X2, tt, NULL);
295:     idiff = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
296:     PetscDSSetObjective(prob, 0, &f0_0_maxwellian_lp);
297:     DMPlexComputeIntegralFEM(plex, X2, tt, NULL);
298:     lpm1 = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
299:   }
300:   PetscPrintf(PETSC_COMM_WORLD, "%s %" PetscInt_FMT ") time=%10.3e n-%d norm electrons/max=%20.13e ions/max=%20.13e\n", "----", stepi, (double)time, (int)ppp, (double)(ediff / lpm0), (double)(idiff / lpm1));
301:   /* view */
302:   VecCopy(X2, X);
303:   VecDestroy(&X2);
304:   if (islast) {
305:     VecDestroy(&rectx->X_0);
306:     rectx->X_0 = NULL;
307:   }
308:   return 0;
309: }

311: static PetscErrorCode EInduction(Vec X, Vec X_t, PetscInt step, PetscReal time, LandauCtx *ctx, PetscReal *a_E)
312: {
313:   REctx      *rectx = (REctx *)ctx->data;
314:   PetscInt    ii;
315:   DM          dm, plex;
316:   PetscScalar tt[LANDAU_MAX_SPECIES], qv0[LANDAU_MAX_SPECIES];
317:   PetscReal   dJ_dt;
318:   PetscDS     prob;

321:   for (ii = 0; ii < ctx->num_species; ii++) qv0[ii] = ctx->charges[ii] * ctx->v_0;
322:   VecGetDM(X, &dm);
323:   DMGetDS(dm, &prob);
324:   DMConvert(dm, DMPLEX, &plex);
325:   /* get d current / dt */
326:   PetscDSSetConstants(prob, ctx->num_species, qv0);
327:   PetscDSSetObjective(prob, 0, &f0_jz_sum);
329:   DMPlexComputeIntegralFEM(plex, X_t, tt, NULL);
330:   dJ_dt = -ctx->n_0 * PetscRealPart(tt[0]) / ctx->t_0;
331:   /* E induction */
332:   *a_E = -rectx->L * dJ_dt + rectx->Ez_initial;
333:   DMDestroy(&plex);
334:   return 0;
335: }

337: static PetscErrorCode EConstant(Vec X, Vec X_t, PetscInt step, PetscReal time, LandauCtx *ctx, PetscReal *a_E)
338: {
340:   *a_E = ctx->Ez;
341:   return 0;
342: }

344: static PetscErrorCode ENone(Vec X, Vec X_t, PetscInt step, PetscReal time, LandauCtx *ctx, PetscReal *a_E)
345: {
347:   *a_E = 0;
348:   return 0;
349: }

351: /* ------------------------------------------------------------------- */
352: /*
353:    FormSource - Evaluates source terms F(t).

355:    Input Parameters:
356: .  ts - the TS context
357: .  time -
358: .  X_dummmy - input vector
359: .  dummy - optional user-defined context, as set by SNESSetFunction()

361:    Output Parameter:
362: .  F - function vector
363:  */
364: static PetscErrorCode FormSource(TS ts, PetscReal ftime, Vec X_dummmy, Vec F, void *dummy)
365: {
366:   PetscReal  new_imp_rate;
367:   LandauCtx *ctx;
368:   DM         pack;
369:   REctx     *rectx;

372:   TSGetDM(ts, &pack);
373:   DMGetApplicationContext(pack, &ctx);
374:   rectx = (REctx *)ctx->data;
375:   /* check for impurities */
376:   rectx->impuritySrcRate(ftime, &new_imp_rate, ctx);
377:   if (new_imp_rate != 0) {
378:     if (new_imp_rate != rectx->current_rate) {
379:       PetscInt  ii;
380:       PetscReal dne_dt, dni_dt, tilda_ns[LANDAU_MAX_SPECIES], temps[LANDAU_MAX_SPECIES];
381:       Vec       globFarray[LANDAU_MAX_GRIDS * LANDAU_MAX_BATCH_SZ];
382:       rectx->current_rate = new_imp_rate;
383:       for (ii = 1; ii < LANDAU_MAX_SPECIES; ii++) tilda_ns[ii] = 0;
384:       for (ii = 1; ii < LANDAU_MAX_SPECIES; ii++) temps[ii] = 1;
385:       dni_dt                   = new_imp_rate /* *ctx->t_0 */; /* fully ionized immediately, no normalize, stay in non-dim */
386:       dne_dt                   = new_imp_rate * rectx->Ne_ion /* *ctx->t_0 */;
387:       tilda_ns[0]              = dne_dt;
388:       tilda_ns[rectx->imp_idx] = dni_dt;
389:       temps[0]                 = rectx->T_cold;
390:       temps[rectx->imp_idx]    = rectx->T_cold;
391:       PetscInfo(ctx->plex[0], "\tHave new_imp_rate= %10.3e time= %10.3e de/dt= %10.3e di/dt= %10.3e ***\n", (double)new_imp_rate, (double)ftime, (double)dne_dt, (double)dni_dt);
392:       DMCompositeGetAccessArray(pack, F, ctx->num_grids * ctx->batch_sz, NULL, globFarray);
393:       for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
394:         /* add it */
395:         DMPlexLandauAddMaxwellians(ctx->plex[grid], globFarray[LAND_PACK_IDX(0, grid)], ftime, temps, tilda_ns, grid, 0, 1, ctx);
396:       }
397:       // Does DMCompositeRestoreAccessArray copy the data back? (no)
398:       DMCompositeRestoreAccessArray(pack, F, ctx->num_grids * ctx->batch_sz, NULL, globFarray);
399:     }
400:   } else {
401:     VecZeroEntries(F);
402:     rectx->current_rate = 0;
403:   }
404:   return 0;
405: }

407: PetscErrorCode Monitor(TS ts, PetscInt stepi, PetscReal time, Vec X, void *actx)
408: {
409:   LandauCtx        *ctx   = (LandauCtx *)actx; /* user-defined application context */
410:   REctx            *rectx = (REctx *)ctx->data;
411:   DM                pack  = NULL;
412:   Vec               globXArray[LANDAU_MAX_GRIDS * LANDAU_MAX_BATCH_SZ];
413:   TSConvergedReason reason;

416:   TSGetConvergedReason(ts, &reason);
417:   if (rectx->grid_view_idx != -1 || (reason && ctx->verbose > 3)) {
418:     VecGetDM(X, &pack);
419:     DMCompositeGetAccessArray(pack, X, ctx->num_grids * ctx->batch_sz, NULL, globXArray);
420:   }
421:   if (stepi > rectx->plotStep && rectx->plotting) {
422:     rectx->plotting = PETSC_FALSE; /* was doing diagnostics, now done */
423:     rectx->plotIdx++;
424:   }
425:   /* view */
426:   if (time / rectx->plotDt >= (PetscReal)rectx->plotIdx || reason) {
427:     if ((reason || stepi == 0 || rectx->plotIdx % rectx->print_period == 0) && ctx->verbose > 1) {
428:       /* print norms */
429:       DMPlexLandauPrintNorms(X, stepi);
430:     }
431:     if (!rectx->plotting) { /* first step of possible backtracks */
432:       rectx->plotting = PETSC_TRUE;
433:       /* diagnostics + change E field with Sptizer (not just a monitor) */
434:       rectx->test(ts, X, stepi, time, reason ? PETSC_TRUE : PETSC_FALSE, ctx, rectx);
435:     } else {
436:       PetscPrintf(PETSC_COMM_WORLD, "\t\t ERROR SKIP test spit ------\n");
437:       rectx->plotting = PETSC_TRUE;
438:     }
439:     if (rectx->grid_view_idx != -1) {
440:       PetscObjectSetName((PetscObject)globXArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], rectx->grid_view_idx == 0 ? "ue" : "ui");
441:       /* view, overwrite step when back tracked */
442:       DMSetOutputSequenceNumber(ctx->plex[rectx->grid_view_idx], rectx->plotIdx, time * ctx->t_0);
443:       VecViewFromOptions(globXArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], NULL, "-ex2_vec_view");
444:     }
445:     rectx->plotStep = stepi;
446:   } else {
447:     if (rectx->plotting) PetscPrintf(PETSC_COMM_WORLD, " ERROR rectx->plotting=%s step %" PetscInt_FMT "\n", PetscBools[rectx->plotting], stepi);
448:     /* diagnostics + change E field with Sptizer (not just a monitor) - can we lag this? */
449:     rectx->test(ts, X, stepi, time, reason ? PETSC_TRUE : PETSC_FALSE, ctx, rectx);
450:   }
451:   /* parallel check that only works of all batches are identical */
452:   if (reason && ctx->verbose > 3) {
453:     PetscReal   val, rval;
454:     PetscMPIInt rank;
455:     MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
456:     for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
457:       PetscInt nerrors = 0;
458:       for (PetscInt i = 0; i < ctx->batch_sz; i++) {
459:         VecNorm(globXArray[LAND_PACK_IDX(i, grid)], NORM_2, &val);
460:         if (i == 0) rval = val;
461:         else if ((val = PetscAbs(val - rval) / rval) > 1000 * PETSC_MACHINE_EPSILON) {
462:           PetscPrintf(PETSC_COMM_SELF, " [%d] Warning %" PetscInt_FMT ".%" PetscInt_FMT ") diff = %2.15e\n", rank, grid, i, (double)val);
463:           nerrors++;
464:         }
465:       }
466:       if (nerrors) {
467:         PetscPrintf(PETSC_COMM_SELF, " ***** [%d] ERROR max %" PetscInt_FMT " errors\n", rank, nerrors);
468:       } else {
469:         PetscPrintf(PETSC_COMM_WORLD, "[%d] %" PetscInt_FMT ") batch consistency check OK\n", rank, grid);
470:       }
471:     }
472:   }
473:   rectx->idx = 0;
474:   if (rectx->grid_view_idx != -1 || (reason && ctx->verbose > 3)) DMCompositeRestoreAccessArray(pack, X, ctx->num_grids * ctx->batch_sz, NULL, globXArray);
475:   return 0;
476: }

478: PetscErrorCode PreStep(TS ts)
479: {
480:   LandauCtx *ctx;
481:   REctx     *rectx;
482:   DM         dm;
483:   PetscInt   stepi;
484:   PetscReal  time;
485:   Vec        X;

488:   /* not used */
489:   TSGetDM(ts, &dm);
490:   TSGetTime(ts, &time);
491:   TSGetSolution(ts, &X);
492:   DMGetApplicationContext(dm, &ctx);
493:   rectx = (REctx *)ctx->data;
494:   TSGetStepNumber(ts, &stepi);
495:   /* update E */
496:   rectx->E(X, NULL, stepi, time, ctx, &ctx->Ez);
497:   return 0;
498: }

500: /* model for source of non-ionized impurities, profile provided by model, in du/dt form in normalized units (tricky because n_0 is normalized with electrons) */
501: static PetscErrorCode stepSrc(PetscReal time, PetscReal *rho, LandauCtx *ctx)
502: {
503:   REctx *rectx = (REctx *)ctx->data;

506:   if (time >= rectx->pulse_start) *rho = rectx->pulse_rate;
507:   else *rho = 0.;
508:   return 0;
509: }
510: static PetscErrorCode zeroSrc(PetscReal time, PetscReal *rho, LandauCtx *ctx)
511: {
513:   *rho = 0.;
514:   return 0;
515: }
516: static PetscErrorCode pulseSrc(PetscReal time, PetscReal *rho, LandauCtx *ctx)
517: {
518:   REctx *rectx = (REctx *)ctx->data;

522:   if (time < rectx->pulse_start || time > rectx->pulse_start + 3 * rectx->pulse_width) *rho = 0;
523:   else {
524:     double x = PetscSinReal((time - rectx->pulse_start) / (3 * rectx->pulse_width) * 2 * PETSC_PI - PETSC_PI / 2) + 1; /* 0:2, integrates to 1.0 */
525:     *rho     = rectx->pulse_rate * x / (3 * rectx->pulse_width);
526:     if (!rectx->use_spitzer_eta) rectx->use_spitzer_eta = PETSC_TRUE; /* use it next time */
527:   }
528:   return 0;
529: }

533: static PetscErrorCode ProcessREOptions(REctx *rectx, const LandauCtx *ctx, DM dm, const char prefix[])
534: {
535:   PetscFunctionList plist = NULL, testlist = NULL, elist = NULL;
536:   char              pname[256], testname[256], ename[256];
537:   DM                dm_dummy;
538:   PetscBool         Connor_E = PETSC_FALSE;

541:   DMCreate(PETSC_COMM_WORLD, &dm_dummy);
542:   rectx->Ne_ion          = 1;    /* number of electrons given up by impurity ion */
543:   rectx->T_cold          = .005; /* kev */
544:   rectx->ion_potential   = 15;   /* ev */
545:   rectx->L               = 2;
546:   rectx->X_0             = NULL;
547:   rectx->imp_idx         = ctx->num_species - 1; /* default ionized impurity as last one */
548:   rectx->pulse_start     = PETSC_MAX_REAL;
549:   rectx->pulse_width     = 1;
550:   rectx->plotStep        = PETSC_MAX_INT;
551:   rectx->pulse_rate      = 1.e-1;
552:   rectx->current_rate    = 0;
553:   rectx->plotIdx         = 0;
554:   rectx->j               = 0;
555:   rectx->plotDt          = 1.0;
556:   rectx->plotting        = PETSC_FALSE;
557:   rectx->use_spitzer_eta = PETSC_FALSE;
558:   rectx->idx             = 0;
559:   rectx->print_period    = 10;
560:   rectx->grid_view_idx   = -1; // do not get if not needed
561:   /* Register the available impurity sources */
562:   PetscFunctionListAdd(&plist, "step", &stepSrc);
563:   PetscFunctionListAdd(&plist, "none", &zeroSrc);
564:   PetscFunctionListAdd(&plist, "pulse", &pulseSrc);
565:   PetscStrcpy(pname, "none");
566:   PetscFunctionListAdd(&testlist, "none", &testNone);
567:   PetscFunctionListAdd(&testlist, "spitzer", &testSpitzer);
568:   PetscFunctionListAdd(&testlist, "stable", &testStable);
569:   PetscStrcpy(testname, "none");
570:   PetscFunctionListAdd(&elist, "none", &ENone);
571:   PetscFunctionListAdd(&elist, "induction", &EInduction);
572:   PetscFunctionListAdd(&elist, "constant", &EConstant);
573:   PetscStrcpy(ename, "constant");

575:   PetscOptionsBegin(PETSC_COMM_SELF, prefix, "Options for Runaway/seed electron model", "none");
576:   PetscOptionsReal("-ex2_plot_dt", "Plotting interval", "ex2.c", rectx->plotDt, &rectx->plotDt, NULL);
577:   if (rectx->plotDt < 0) rectx->plotDt = 1e30;
578:   if (rectx->plotDt == 0) rectx->plotDt = 1e-30;
579:   PetscOptionsInt("-ex2_print_period", "Plotting interval", "ex2.c", rectx->print_period, &rectx->print_period, NULL);
580:   PetscOptionsInt("-ex2_grid_view_idx", "grid_view_idx", "ex2.c", rectx->grid_view_idx, &rectx->grid_view_idx, NULL);
582:   PetscOptionsFList("-ex2_impurity_source_type", "Name of impurity source to run", "", plist, pname, pname, sizeof(pname), NULL);
583:   PetscOptionsFList("-ex2_test_type", "Name of test to run", "", testlist, testname, testname, sizeof(testname), NULL);
584:   PetscOptionsInt("-ex2_impurity_index", "index of sink for impurities", "none", rectx->imp_idx, &rectx->imp_idx, NULL);
586:   PetscOptionsFList("-ex2_e_field_type", "Electric field type", "", elist, ename, ename, sizeof(ename), NULL);
587:   rectx->Ne_ion = -ctx->charges[rectx->imp_idx] / ctx->charges[0];
588:   PetscOptionsReal("-ex2_t_cold", "Temperature of cold electron and ions after ionization in keV", "none", rectx->T_cold, &rectx->T_cold, NULL);
589:   PetscOptionsReal("-ex2_pulse_start_time", "Time at which pulse happens for 'pulse' source", "none", rectx->pulse_start, &rectx->pulse_start, NULL);
590:   PetscOptionsReal("-ex2_pulse_width_time", "Width of pulse 'pulse' source", "none", rectx->pulse_width, &rectx->pulse_width, NULL);
591:   PetscOptionsReal("-ex2_pulse_rate", "Number density of pulse for 'pulse' source", "none", rectx->pulse_rate, &rectx->pulse_rate, NULL);
592:   rectx->T_cold *= 1.16e7; /* convert to Kelvin */
593:   PetscOptionsReal("-ex2_ion_potential", "Potential to ionize impurity (should be array) in ev", "none", rectx->ion_potential, &rectx->ion_potential, NULL);
594:   PetscOptionsReal("-ex2_inductance", "Inductance E feild", "none", rectx->L, &rectx->L, NULL);
595:   PetscOptionsBool("-ex2_connor_e_field_units", "Scale Ex but Connor-Hastie E_c", "none", Connor_E, &Connor_E, NULL);
596:   PetscInfo(dm_dummy, "Num electrons from ions=%g, T_cold=%10.3e, ion potential=%10.3e, E_z=%10.3e v_0=%10.3e\n", (double)rectx->Ne_ion, (double)rectx->T_cold, (double)rectx->ion_potential, (double)ctx->Ez, (double)ctx->v_0);
597:   PetscOptionsEnd();
598:   /* get impurity source rate function */
599:   PetscFunctionListFind(plist, pname, &rectx->impuritySrcRate);
601:   PetscFunctionListFind(testlist, testname, &rectx->test);
603:   PetscFunctionListFind(elist, ename, &rectx->E);
605:   PetscFunctionListDestroy(&plist);
606:   PetscFunctionListDestroy(&testlist);
607:   PetscFunctionListDestroy(&elist);

609:   /* convert E from Connor-Hastie E_c units to real if doing Spitzer E */
610:   if (Connor_E) {
611:     PetscReal E = ctx->Ez, Tev = ctx->thermal_temps[0] * 8.621738e-5, n = ctx->n_0 * ctx->n[0];
612:     CalculateE(Tev, n, ctx->lnLam, ctx->epsilon0, &E);
613:     ((LandauCtx *)ctx)->Ez *= E;
614:   }
615:   DMDestroy(&dm_dummy);
616:   return 0;
617: }

619: int main(int argc, char **argv)
620: {
621:   DM         pack;
622:   Vec        X;
623:   PetscInt   dim = 2, nDMs;
624:   TS         ts;
625:   Mat        J;
626:   PetscDS    prob;
627:   LandauCtx *ctx;
628:   REctx     *rectx;
629: #if defined PETSC_USE_LOG
630:   PetscLogStage stage;
631: #endif
632:   PetscMPIInt rank;
633: #if defined(PETSC_HAVE_THREADSAFETY)
634:   double starttime, endtime;
635: #endif
637:   PetscInitialize(&argc, &argv, NULL, help);
638:   MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
639:   if (rank) { /* turn off output stuff for duplicate runs */
640:     PetscOptionsClearValue(NULL, "-ex2_dm_view");
641:     PetscOptionsClearValue(NULL, "-ex2_vec_view");
642:     PetscOptionsClearValue(NULL, "-ex2_vec_view_init");
643:     PetscOptionsClearValue(NULL, "-ex2_dm_view_init");
644:     PetscOptionsClearValue(NULL, "-info"); /* this does not work */
645:   }
646:   PetscOptionsGetInt(NULL, NULL, "-dim", &dim, NULL);
647:   /* Create a mesh */
648:   DMPlexLandauCreateVelocitySpace(PETSC_COMM_WORLD, dim, "", &X, &J, &pack);
649:   DMCompositeGetNumberDM(pack, &nDMs);
650:   PetscObjectSetName((PetscObject)J, "Jacobian");
651:   PetscObjectSetName((PetscObject)X, "f");
652:   DMGetApplicationContext(pack, &ctx);
653:   DMSetUp(pack);
654:   /* context */
655:   PetscNew(&rectx);
656:   ctx->data = rectx;
657:   ProcessREOptions(rectx, ctx, pack, "");
658:   DMGetDS(pack, &prob);
659:   if (rectx->grid_view_idx != -1) {
660:     Vec *XsubArray = NULL;
661:     PetscMalloc(sizeof(*XsubArray) * nDMs, &XsubArray);
662:     DMCompositeGetAccessArray(pack, X, nDMs, NULL, XsubArray); // read only
663:     PetscObjectSetName((PetscObject)XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], rectx->grid_view_idx == 0 ? "ue" : "ui");
664:     DMSetOutputSequenceNumber(ctx->plex[rectx->grid_view_idx], 0, 0.0);
665:     DMViewFromOptions(ctx->plex[rectx->grid_view_idx], NULL, "-ex2_dm_view");
666:     DMViewFromOptions(ctx->plex[rectx->grid_view_idx], NULL, "-ex2_dm_view_init");
667:     VecViewFromOptions(XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], NULL, "-ex2_vec_view");      // initial condition (monitor plots after step)
668:     VecViewFromOptions(XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], NULL, "-ex2_vec_view_init"); // initial condition (monitor plots after step)
669:     DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, XsubArray);                                                       // read only
670:     PetscFree(XsubArray);
671:   }
672:   /* Create timestepping solver context */
673:   TSCreate(PETSC_COMM_SELF, &ts);
674:   TSSetDM(ts, pack);
675:   TSSetIFunction(ts, NULL, DMPlexLandauIFunction, NULL);
676:   TSSetIJacobian(ts, J, J, DMPlexLandauIJacobian, NULL);
677:   TSSetRHSFunction(ts, NULL, FormSource, NULL);
678:   TSSetFromOptions(ts);
679:   TSSetSolution(ts, X);
680:   TSSetApplicationContext(ts, ctx);
681:   TSMonitorSet(ts, Monitor, ctx, NULL);
682:   TSSetPreStep(ts, PreStep);
683:   rectx->Ez_initial = ctx->Ez; /* cache for induction caclulation - applied E field */
684:   if (1) {                     /* warm up an test just DMPlexLandauIJacobian */
685:     Vec       vec;
686:     PetscInt  nsteps;
687:     PetscReal dt;
688:     PetscLogStageRegister("Warmup", &stage);
689:     PetscLogStagePush(stage);
690:     VecDuplicate(X, &vec);
691:     VecCopy(X, vec);
692:     TSGetMaxSteps(ts, &nsteps);
693:     TSGetTimeStep(ts, &dt);
694:     TSSetMaxSteps(ts, 1);
695:     TSSolve(ts, X);
696:     TSSetMaxSteps(ts, nsteps);
697:     TSSetStepNumber(ts, 0);
698:     TSSetTime(ts, 0);
699:     TSSetTimeStep(ts, dt);
700:     rectx->plotIdx  = 0;
701:     rectx->plotting = PETSC_FALSE;
702:     PetscLogStagePop();
703:     VecCopy(vec, X);
704:     VecDestroy(&vec);
705:     PetscObjectStateIncrease((PetscObject)ctx->J);
706:   }
707:   /* go */
708:   PetscLogStageRegister("Solve", &stage);
709:   ctx->stage = 0; // lets not use this stage
710: #if defined(PETSC_HAVE_THREADSAFETY)
711:   ctx->stage = 1; // not set with thread safety
712: #endif
713:   //TSSetSolution(ts,X);
714:   PetscLogStagePush(stage);
715: #if defined(PETSC_HAVE_THREADSAFETY)
716:   starttime = MPI_Wtime();
717: #endif
718: #if defined(PETSC_HAVE_CUDA_NVTX)
719:   nvtxRangePushA("ex2-TSSolve-warm");
720: #endif
721:   TSSolve(ts, X);
722: #if defined(PETSC_HAVE_CUDA_NVTX)
723:   nvtxRangePop();
724: #endif
725:   PetscLogStagePop();
726: #if defined(PETSC_HAVE_THREADSAFETY)
727:   endtime = MPI_Wtime();
728:   ctx->times[LANDAU_EX2_TSSOLVE] += (endtime - starttime);
729: #endif
730:   /* clean up */
731:   DMPlexLandauDestroyVelocitySpace(&pack);
732:   TSDestroy(&ts);
733:   VecDestroy(&X);
734:   PetscFree(rectx);
735:   PetscFinalize();
736:   return 0;
737: }

739: /*TEST

741:   testset:
742:     requires: p4est !complex double defined(PETSC_USE_DMLANDAU_2D)
743:     output_file: output/ex2_0.out
744:     args: -dm_landau_num_species_grid 1,1 -dm_landau_Ez 0 -petscspace_degree 3 -petscspace_poly_tensor 1 -dm_landau_type p4est -dm_landau_ion_masses 2 -dm_landau_ion_charges 1 -dm_landau_thermal_temps 5,5 -dm_landau_n 2,2 -dm_landau_n_0 5e19 -ts_monitor -snes_rtol 1.e-10 -snes_stol 1.e-14 -snes_monitor -snes_converged_reason -snes_max_it 10 -ts_type arkimex -ts_arkimex_type 1bee -ts_max_snes_failures -1 -ts_rtol 1e-3 -ts_dt 1.e-1 -ts_max_time 1 -ts_adapt_clip .5,1.25 -ts_max_steps 2 -ts_adapt_scale_solve_failed 0.75 -ts_adapt_time_step_increase_delay 5 -dm_landau_amr_levels_max 2,2 -ex2_impurity_source_type pulse -ex2_pulse_start_time 1e-1 -ex2_pulse_width_time 10 -ex2_pulse_rate 1e-2 -ex2_t_cold .05 -ex2_plot_dt 1e-1 -dm_refine 0 -dm_landau_gpu_assembly true -dm_landau_batch_size 2 -dm_landau_verbose 2
745:     test:
746:       suffix: cpu
747:       args: -dm_landau_device_type cpu -ksp_type bicg -pc_type jacobi
748:     test:
749:       suffix: kokkos
750:       requires: kokkos_kernels
751:       args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -ksp_type bicg -pc_type jacobi
752:     test:
753:       suffix: cuda
754:       requires: cuda
755:       args: -dm_landau_device_type cuda -dm_mat_type aijcusparse -dm_vec_type cuda -mat_cusparse_use_cpu_solve -ksp_type bicg -pc_type jacobi
756:     test:
757:       suffix: kokkos_batch
758:       requires: kokkos_kernels
759:       args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -ksp_type preonly -pc_type bjkokkos -pc_bjkokkos_ksp_type bicg -pc_bjkokkos_pc_type jacobi
760:     test:
761:       suffix: kokkos_batch_coo
762:       requires: kokkos_kernels
763:       args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -ksp_type preonly -pc_type bjkokkos -pc_bjkokkos_ksp_type bicg -pc_bjkokkos_pc_type jacobi -dm_landau_coo_assembly

765: TEST*/