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:   #if PETSC_PKG_CUDA_VERSION_GE(10, 0, 0)
     12:     #include <nvtx3/nvToolsExt.h>
     13:   #else
     14:     #include <nvToolsExt.h>
     15:   #endif
     16: #endif

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

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

     47: #define RE_CUT 3.
     48: /* < v, u_re * v * q > */
     49: 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)
     50: {
     51:   PetscReal n_e = PetscRealPart(u[0]);
     52:   if (dim == 2) {
     53:     if (x[1] > RE_CUT || x[1] < -RE_CUT) {                    /* simply a cutoff for REs. v_|| > 3 v(T_e) */
     54:       *f0 = n_e * 2. * PETSC_PI * x[0] * x[1] * constants[0]; /* n * r * v_|| * q */
     55:     } else {
     56:       *f0 = 0;
     57:     }
     58:   } else {
     59:     if (x[2] > RE_CUT || x[2] < -RE_CUT) { /* simply a cutoff for REs. v_|| > 3 v(T_e) */
     60:       *f0 = n_e * x[2] * constants[0];
     61:     } else {
     62:       *f0 = 0;
     63:     }
     64:   }
     65: }

     67: /* sum < v, u*v*q > */
     68: 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)
     69: {
     70:   PetscInt ii;
     71:   f0[0] = 0;
     72:   if (dim == 2) {
     73:     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 */
     74:   } else {
     75:     for (ii = 0; ii < Nf; ii++) f0[0] += u[ii] * x[2] * q[ii]; /* n * v_|| * q  * v_0 */
     76:   }
     77: }

     79: /* < v, n_e > */
     80: 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)
     81: {
     82:   PetscInt ii = (PetscInt)PetscRealPart(constants[0]);
     83:   if (dim == 2) f0[0] = 2. * PETSC_PI * x[0] * u[ii];
     84:   else f0[0] = u[ii];
     85: }

     87: /* < v, n_e v_|| > */
     88: 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)
     89: {
     90:   PetscInt ii = (PetscInt)PetscRealPart(constants[0]);
     91:   if (dim == 2) f0[0] = u[ii] * 2. * PETSC_PI * x[0] * x[1]; /* n r v_|| */
     92:   else f0[0] = u[ii] * x[2];                                 /* n v_|| */
     93: }

     95: /* < v, n_e (v-shift) > */
     96: 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)
     97: {
     98:   PetscReal vz = numConstants > 0 ? PetscRealPart(constants[0]) : 0;
     99:   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 */
    100:   else {
    101:     *f0 = u[0] * PetscSqrtReal(x[0] * x[0] + x[1] * x[1] + (x[2] - vz) * (x[2] - vz)); /* n v */
    102:   }
    103: }

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

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

    133: static PetscReal Spitzer(PetscReal m_e, PetscReal e, PetscReal Z, PetscReal epsilon0, PetscReal lnLam, PetscReal kTe_joules)
    134: {
    135:   PetscReal Fz = (1 + 1.198 * Z + 0.222 * Z * Z) / (1 + 2.966 * Z + 0.753 * Z * Z), eta;
    136:   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);
    137:   return eta;
    138: }

    140: static PetscErrorCode testNone(TS ts, Vec X, PetscInt stepi, PetscReal time, PetscBool islast, LandauCtx *ctx, REctx *rectx)
    141: {
    142:   PetscFunctionBeginUser;
    143:   PetscFunctionReturn(PETSC_SUCCESS);
    144: }

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

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

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

    224:   ratio = E / J / spit_eta;
    225:   if (stepi > 10 && !rectx->use_spitzer_eta && (old_ratio - ratio < 1.e-6)) {
    226:     rectx->pulse_start     = time + 0.98 * dt;
    227:     rectx->use_spitzer_eta = PETSC_TRUE;
    228:   }
    229:   PetscCall(TSGetConvergedReason(ts, &reason));
    230:   PetscCall(TSGetConvergedReason(ts, &reason));
    231:   if (rectx->plotting || stepi == 0 || reason || rectx->pulse_start == time + 0.98 * dt) {
    232:     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,
    233:                           (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));
    234:     PetscCheck(rectx->pulse_start != (time + 0.98 * dt), PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "Spitzer complete ratio=%g", (double)ratio);
    235:   }
    236:   old_ratio = ratio;
    237:   PetscFunctionReturn(PETSC_SUCCESS);
    238: }

    240: static const double ppp = 2;
    241: 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)
    242: {
    243:   LandauCtx      *ctx   = (LandauCtx *)constants;
    244:   REctx          *rectx = (REctx *)ctx->data;
    245:   PetscInt        ii    = rectx->idx, i;
    246:   const PetscReal kT_m  = ctx->k * ctx->thermal_temps[ii] / ctx->masses[ii]; /* kT/m */
    247:   const PetscReal n     = ctx->n[ii];
    248:   PetscReal       diff, f_maxwell, v2 = 0, theta = 2 * kT_m / (ctx->v_0 * ctx->v_0); /* theta = 2kT/mc^2 */
    249:   for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
    250:   f_maxwell = n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
    251:   diff      = 2. * PETSC_PI * x[0] * (PetscRealPart(u[ii]) - f_maxwell);
    252:   f0[0]     = PetscPowReal(diff, ppp);
    253: }
    254: 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)
    255: {
    256:   LandauCtx      *ctx   = (LandauCtx *)constants;
    257:   REctx          *rectx = (REctx *)ctx->data;
    258:   PetscInt        ii    = rectx->idx, i;
    259:   const PetscReal kT_m  = ctx->k * ctx->thermal_temps[ii] / ctx->masses[ii]; /* kT/m */
    260:   const PetscReal n     = ctx->n[ii];
    261:   PetscReal       f_maxwell, v2 = 0, theta = 2 * kT_m / (ctx->v_0 * ctx->v_0); /* theta = 2kT/mc^2 */
    262:   for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
    263:   f_maxwell = 2. * PETSC_PI * x[0] * n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
    264:   f0[0]     = PetscPowReal(f_maxwell, ppp);
    265: }

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

    275:   PetscFunctionBeginUser;
    276:   PetscCall(VecGetDM(X, &dm));
    277:   PetscCall(DMGetDS(plex, &prob));
    278:   PetscCall(VecDuplicate(X, &X2));
    279:   PetscCall(VecCopy(X, X2));
    280:   if (!rectx->X_0) {
    281:     PetscCall(VecDuplicate(X, &rectx->X_0));
    282:     PetscCall(VecCopy(X, rectx->X_0));
    283:   }
    284:   PetscCall(VecAXPY(X, -1.0, rectx->X_0));
    285:   PetscCall(PetscDSSetConstants(prob, sizeof(LandauCtx) / sizeof(PetscScalar), (PetscScalar *)ctx));
    286:   rectx->idx = 0;
    287:   PetscCall(PetscDSSetObjective(prob, 0, &f0_0_diff_lp));
    288:   PetscCall(DMPlexComputeIntegralFEM(plex, X2, tt, NULL));
    289:   ediff = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
    290:   PetscCall(PetscDSSetObjective(prob, 0, &f0_0_maxwellian_lp));
    291:   PetscCall(DMPlexComputeIntegralFEM(plex, X2, tt, NULL));
    292:   lpm0 = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
    293:   if (ctx->num_species > 1) {
    294:     rectx->idx = 1;
    295:     PetscCall(PetscDSSetObjective(prob, 0, &f0_0_diff_lp));
    296:     PetscCall(DMPlexComputeIntegralFEM(plex, X2, tt, NULL));
    297:     idiff = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
    298:     PetscCall(PetscDSSetObjective(prob, 0, &f0_0_maxwellian_lp));
    299:     PetscCall(DMPlexComputeIntegralFEM(plex, X2, tt, NULL));
    300:     lpm1 = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
    301:   }
    302:   PetscCall(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)));
    303:   /* view */
    304:   PetscCall(VecCopy(X2, X));
    305:   PetscCall(VecDestroy(&X2));
    306:   if (islast) {
    307:     PetscCall(VecDestroy(&rectx->X_0));
    308:     rectx->X_0 = NULL;
    309:   }
    310:   PetscFunctionReturn(PETSC_SUCCESS);
    311: }

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

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

    339: static PetscErrorCode EConstant(Vec X, Vec X_t, PetscInt step, PetscReal time, LandauCtx *ctx, PetscReal *a_E)
    340: {
    341:   PetscFunctionBeginUser;
    342:   *a_E = ctx->Ez;
    343:   PetscFunctionReturn(PETSC_SUCCESS);
    344: }

    346: static PetscErrorCode ENone(Vec X, Vec X_t, PetscInt step, PetscReal time, LandauCtx *ctx, PetscReal *a_E)
    347: {
    348:   PetscFunctionBeginUser;
    349:   *a_E = 0;
    350:   PetscFunctionReturn(PETSC_SUCCESS);
    351: }

    353: /* ------------------------------------------------------------------- */
    354: /*
    355:    FormSource - Evaluates source terms F(t).

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

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

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

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

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

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

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

    502: /* 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) */
    503: static PetscErrorCode stepSrc(PetscReal time, PetscReal *rho, LandauCtx *ctx)
    504: {
    505:   REctx *rectx = (REctx *)ctx->data;

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

    522:   PetscFunctionBeginUser;
    523:   PetscCheck(rectx->pulse_start != PETSC_MAX_REAL, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "'-ex2_pulse_start_time X' must be used with '-ex2_impurity_source_type pulse'");
    524:   if (time < rectx->pulse_start || time > rectx->pulse_start + 3 * rectx->pulse_width) *rho = 0;
    525:   else {
    526:     double x = PetscSinReal((time - rectx->pulse_start) / (3 * rectx->pulse_width) * 2 * PETSC_PI - PETSC_PI / 2) + 1; /* 0:2, integrates to 1.0 */
    527:     *rho     = rectx->pulse_rate * x / (3 * rectx->pulse_width);
    528:     if (!rectx->use_spitzer_eta) rectx->use_spitzer_eta = PETSC_TRUE; /* use it next time */
    529:   }
    530:   PetscFunctionReturn(PETSC_SUCCESS);
    531: }

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

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

    577:   PetscOptionsBegin(PETSC_COMM_SELF, prefix, "Options for Runaway/seed electron model", "none");
    578:   PetscCall(PetscOptionsReal("-ex2_plot_dt", "Plotting interval", "ex2.c", rectx->plotDt, &rectx->plotDt, NULL));
    579:   if (rectx->plotDt < 0) rectx->plotDt = 1e30;
    580:   if (rectx->plotDt == 0) rectx->plotDt = 1e-30;
    581:   PetscCall(PetscOptionsInt("-ex2_print_period", "Plotting interval", "ex2.c", rectx->print_period, &rectx->print_period, NULL));
    582:   PetscCall(PetscOptionsInt("-ex2_grid_view_idx", "grid_view_idx", "ex2.c", rectx->grid_view_idx, &rectx->grid_view_idx, NULL));
    583:   PetscCheck(rectx->grid_view_idx < ctx->num_grids || rectx->grid_view_idx == -1, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "rectx->grid_view_idx (%" PetscInt_FMT ") >= ctx->num_grids (%" PetscInt_FMT ")", rectx->imp_idx, ctx->num_grids);
    584:   PetscCall(PetscOptionsFList("-ex2_impurity_source_type", "Name of impurity source to run", "", plist, pname, pname, sizeof(pname), NULL));
    585:   PetscCall(PetscOptionsFList("-ex2_test_type", "Name of test to run", "", testlist, testname, testname, sizeof(testname), NULL));
    586:   PetscCall(PetscOptionsInt("-ex2_impurity_index", "index of sink for impurities", "none", rectx->imp_idx, &rectx->imp_idx, NULL));
    587:   PetscCheck((rectx->imp_idx < ctx->num_species && rectx->imp_idx >= 1) || ctx->num_species <= 1, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "index of sink for impurities ions is out of range (%" PetscInt_FMT "), must be > 0 && < NS", rectx->imp_idx);
    588:   PetscCall(PetscOptionsFList("-ex2_e_field_type", "Electric field type", "", elist, ename, ename, sizeof(ename), NULL));
    589:   rectx->Ne_ion = -ctx->charges[rectx->imp_idx] / ctx->charges[0];
    590:   PetscCall(PetscOptionsReal("-ex2_t_cold", "Temperature of cold electron and ions after ionization in keV", "none", rectx->T_cold, &rectx->T_cold, NULL));
    591:   PetscCall(PetscOptionsReal("-ex2_pulse_start_time", "Time at which pulse happens for 'pulse' source", "none", rectx->pulse_start, &rectx->pulse_start, NULL));
    592:   PetscCall(PetscOptionsReal("-ex2_pulse_width_time", "Width of pulse 'pulse' source", "none", rectx->pulse_width, &rectx->pulse_width, NULL));
    593:   PetscCall(PetscOptionsReal("-ex2_pulse_rate", "Number density of pulse for 'pulse' source", "none", rectx->pulse_rate, &rectx->pulse_rate, NULL));
    594:   rectx->T_cold *= 1.16e7; /* convert to Kelvin */
    595:   PetscCall(PetscOptionsReal("-ex2_ion_potential", "Potential to ionize impurity (should be array) in ev", "none", rectx->ion_potential, &rectx->ion_potential, NULL));
    596:   PetscCall(PetscOptionsReal("-ex2_inductance", "Inductance E field", "none", rectx->L, &rectx->L, NULL));
    597:   PetscCall(PetscOptionsBool("-ex2_connor_e_field_units", "Scale Ex but Connor-Hastie E_c", "none", Connor_E, &Connor_E, NULL));
    598:   PetscCall(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));
    599:   PetscOptionsEnd();
    600:   /* get impurity source rate function */
    601:   PetscCall(PetscFunctionListFind(plist, pname, &rectx->impuritySrcRate));
    602:   PetscCheck(rectx->impuritySrcRate, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "No impurity source function found '%s'", pname);
    603:   PetscCall(PetscFunctionListFind(testlist, testname, &rectx->test));
    604:   PetscCheck(rectx->test, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "No test found '%s'", testname);
    605:   PetscCall(PetscFunctionListFind(elist, ename, &rectx->E));
    606:   PetscCheck(rectx->E, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "No E field function found '%s'", ename);
    607:   PetscCall(PetscFunctionListDestroy(&plist));
    608:   PetscCall(PetscFunctionListDestroy(&testlist));
    609:   PetscCall(PetscFunctionListDestroy(&elist));

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

    621: int main(int argc, char **argv)
    622: {
    623:   DM            pack;
    624:   Vec           X;
    625:   PetscInt      dim = 2, nDMs;
    626:   TS            ts;
    627:   Mat           J;
    628:   PetscDS       prob;
    629:   LandauCtx    *ctx;
    630:   REctx        *rectx;
    631:   PetscMPIInt   rank;
    632:   PetscLogStage stage;

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

    726: /*TEST

    728:   testset:
    729:     requires: p4est !complex double defined(PETSC_USE_DMLANDAU_2D)
    730:     output_file: output/ex2_0.out
    731:     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-9 -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 unlimited -ts_rtol 1e-3 -ts_dt 1.e-2 -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 -dm_landau_amr_re_levels 2 -dm_landau_re_radius 0 -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 -dm_landau_domain_radius 5.,5.
    732:     test:
    733:       suffix: cpu
    734:       args: -dm_landau_device_type cpu -ksp_type bicg -pc_type jacobi
    735:     test:
    736:       suffix: kokkos
    737:       requires: kokkos_kernels !defined(PETSC_HAVE_CUDA_CLANG)
    738:       args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -ksp_type bicg -pc_type jacobi
    739:     test:
    740:       suffix: kokkos_batch
    741:       requires: kokkos_kernels
    742:       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
    743:     test:
    744:       suffix: kokkos_batch_tfqmr
    745:       requires: kokkos_kernels !defined(PETSC_HAVE_CUDA_CLANG)
    746:       args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -ksp_type preonly -pc_type bjkokkos -pc_bjkokkos_ksp_type tfqmr -pc_bjkokkos_pc_type jacobi

    748:   test:
    749:     requires: !complex double defined(PETSC_USE_DMLANDAU_2D) !cuda
    750:     suffix: single
    751:     nsize: 1
    752:     args: -dm_refine 2 -dm_landau_num_species_grid 1 -dm_landau_thermal_temps 1 -dm_landau_electron_shift 1.25 -petscspace_degree 3 -snes_converged_reason -ts_type beuler -ts_dt .1 -ex2_plot_dt .1 -ts_max_steps 1 -ex2_grid_view_idx 0 -ex2_dm_view -snes_rtol 1.e-13 -snes_stol 1.e-13 -dm_landau_verbose 2 -ex2_print_period 1 -ksp_type preonly -pc_type lu -dm_landau_device_type cpu -dm_landau_use_relativistic_corrections

    754:   testset:
    755:     requires: !complex double defined(PETSC_USE_DMLANDAU_2D)
    756:     nsize: 1
    757:     output_file: output/ex2_simplex.out
    758:     args: -dim 2 -dm_landau_num_species_grid 1,1 -petscspace_degree 2 -dm_landau_simplex -dm_landau_ion_masses 2 -dm_landau_ion_charges 1 -dm_landau_thermal_temps 2,1 -dm_landau_n 1,1 -snes_rtol 1e-15 -snes_stol 1e-15 -snes_monitor -ts_type beuler -snes_converged_reason -ts_exact_final_time stepover -ts_dt .1 -ts_max_steps 1 -ts_max_snes_failures unlimited -ksp_type preonly -pc_type lu -dm_landau_verbose 2 -ex2_grid_view_idx 0 -ex2_dm_view -dm_refine 1 -ksp_type bicg -pc_type jacobi
    759:     test:
    760:       suffix: simplex
    761:       args: -dm_landau_device_type cpu
    762:     test:
    763:       suffix: simplexkokkos
    764:       requires: kokkos_kernels !defined(PETSC_HAVE_CUDA_CLANG) !sycl
    765:       args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos

    767:   test:
    768:     requires: double !defined(PETSC_USE_DMLANDAU_2D)
    769:     suffix: sphere_3d
    770:     nsize: 1
    771:     args: -dim 3 -dm_landau_thermal_temps 2 -petscspace_degree 2 -ts_type beuler -ts_dt .1 -ts_max_steps 1 -dm_landau_verbose 2 -ksp_type preonly -pc_type lu -dm_landau_device_type cpu -snes_rtol 1.e-14 -snes_stol 1.e-14 -snes_converged_reason -dm_landau_sphere -dm_landau_sphere_inner_radius_90degree_scale .55 -dm_landau_sphere_inner_radius_45degree_scale .5

    773: TEST*/