Actual source code: ex1.c

      1: static char help[] = "Landau collision operator with anisotropic thermalization verification test as per Hager et al.\n 'A fully non-linear multi-species Fokker-Planck-Landau collision operator for simulation of fusion plasma', and "
      2:                      "published as 'A performance portable, fully implicit Landau collision operator with batched linear solvers' https://arxiv.org/abs/2209.03228\n\n";

      4: #include <petscts.h>
      5: #include <petsclandau.h>
      6: #include <petscdmcomposite.h>
      7: #include <petscds.h>

      9: /*
     10:  call back method for DMPlexLandauAccess:

     12: Input Parameters:
     13:  .   dm - a DM for this field
     14:  -   local_field - the local index in the grid for this field
     15:  .   grid - the grid index
     16:  +   b_id - the batch index
     17:  -   vctx - a user context

     19:  Input/Output Parameter:
     20:  .   x - Vector to data to

     22:  */
     23: PetscErrorCode landau_field_print_access_callback(DM dm, Vec x, PetscInt local_field, PetscInt grid, PetscInt b_id, void *vctx)
     24: {
     25:   LandauCtx  *ctx;
     26:   PetscScalar val;
     27:   PetscInt    species;

     29:   PetscFunctionBegin;
     30:   PetscCall(DMGetApplicationContext(dm, &ctx));
     31:   species = ctx->species_offset[grid] + local_field;
     32:   val     = (PetscScalar)(LAND_PACK_IDX(b_id, grid) + (species + 1) * 10);
     33:   PetscCall(VecSet(x, val));
     34:   PetscCall(PetscInfo(dm, "DMPlexLandauAccess user 'add' method to grid %" PetscInt_FMT ", batch %" PetscInt_FMT " and local field %" PetscInt_FMT " with %" PetscInt_FMT " grids\n", grid, b_id, local_field, ctx->num_grids));
     35:   PetscFunctionReturn(PETSC_SUCCESS);
     36: }

     38: static const PetscReal alphai   = 1 / 1.3;
     39: static const PetscReal kev_joul = 6.241506479963235e+15; /* 1/1000e */

     41: // constants: [index of (anisotropic) direction of source, z x[1] shift
     42: /* < v, n_s v_|| > */
     43: 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)
     44: {
     45:   if (dim == 2) f0[0] = u[0] * (2. * PETSC_PI * x[0]) * x[1]; /* n r v_|| */
     46:   else f0[0] = u[0] * x[2];
     47: }
     48: /* < v, n (v-shift)^2 > */
     49: static void f0_v2_par_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)
     50: {
     51:   PetscReal vz = PetscRealPart(constants[0]);
     52:   if (dim == 2) *f0 = u[0] * (2. * PETSC_PI * x[0]) * (x[1] - vz) * (x[1] - vz); /* n r v^2_par|perp */
     53:   else *f0 = u[0] * (x[2] - vz) * (x[2] - vz);
     54: }
     55: static void f0_v2_perp(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)
     56: {
     57:   if (dim == 2) *f0 = u[0] * (2. * PETSC_PI * x[0]) * x[0] * x[0]; /* n r v^2_perp */
     58:   else *f0 = u[0] * (x[0] * x[0] + x[1] * x[1]);
     59: }
     60: /* < v, n_e > */
     61: 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)
     62: {
     63:   if (dim == 2) f0[0] = 2. * PETSC_PI * x[0] * u[0];
     64:   else f0[0] = u[0];
     65: }
     66: static void f0_v2_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)
     67: {
     68:   PetscReal vz = PetscRealPart(constants[0]);
     69:   if (dim == 2) f0[0] = u[0] * (2. * PETSC_PI * x[0]) * (x[0] * x[0] + (x[1] - vz) * (x[1] - vz));
     70:   else f0[0] = u[0] * (x[0] * x[0] + x[1] * x[1] + (x[2] - vz) * (x[2] - vz));
     71: }
     72: /* Define a Maxwellian function for testing out the operator. */
     73: typedef struct {
     74:   PetscReal v_0;
     75:   PetscReal kT_m;
     76:   PetscReal n;
     77:   PetscReal shift;
     78:   PetscInt  species;
     79: } MaxwellianCtx;

     81: static PetscErrorCode maxwellian(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf_dummy, PetscScalar *u, void *actx)
     82: {
     83:   MaxwellianCtx *mctx  = (MaxwellianCtx *)actx;
     84:   PetscReal      theta = 2 * mctx->kT_m / (mctx->v_0 * mctx->v_0); /* theta = 2kT/mc^2 */

     86:   PetscFunctionBegin;
     87:   /* evaluate the shifted Maxwellian */
     88:   if (dim == 2) u[0] += alphai * mctx->n * PetscPowReal(PETSC_PI * theta, -1.5) * PetscExpReal(-(alphai * x[0] * x[0] + (x[1] - mctx->shift) * (x[1] - mctx->shift)) / theta);
     89:   else u[0] += alphai * mctx->n * PetscPowReal(PETSC_PI * theta, -1.5) * PetscExpReal(-(alphai * (x[0] * x[0] + x[1] * x[1]) + (x[2] - mctx->shift) * (x[2] - mctx->shift)) / theta);
     90:   PetscFunctionReturn(PETSC_SUCCESS);
     91: }

     93: static PetscErrorCode SetMaxwellians(DM dm, Vec X, PetscReal time, PetscReal temps[], PetscReal ns[], PetscInt grid, PetscReal shifts[], LandauCtx *ctx)
     94: {
     95:   PetscErrorCode (*initu[LANDAU_MAX_SPECIES])(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *);
     96:   MaxwellianCtx *mctxs[LANDAU_MAX_SPECIES], data[LANDAU_MAX_SPECIES];

     98:   PetscFunctionBegin;
     99:   if (!ctx) PetscCall(DMGetApplicationContext(dm, &ctx));
    100:   for (PetscInt ii = ctx->species_offset[grid], i0 = 0; ii < ctx->species_offset[grid + 1]; ii++, i0++) {
    101:     mctxs[i0]        = &data[i0];
    102:     data[i0].v_0     = ctx->v_0;                             // v_0 same for all grids
    103:     data[i0].kT_m    = ctx->k * temps[ii] / ctx->masses[ii]; /* kT/m = v_th ^ 2*/
    104:     data[i0].n       = ns[ii];
    105:     initu[i0]        = maxwellian;
    106:     data[i0].shift   = 0;
    107:     data[i0].species = ii;
    108:   }
    109:   if (1) {
    110:     data[0].shift = (PetscReal)-PetscSign(ctx->charges[ctx->species_offset[grid]]) * ctx->electronShift * ctx->m_0 / ctx->masses[ctx->species_offset[grid]];
    111:   } else {
    112:     shifts[0]     = 0.5 * PetscSqrtReal(ctx->masses[0] / ctx->masses[1]);
    113:     shifts[1]     = 50 * (ctx->masses[0] / ctx->masses[1]);
    114:     data[0].shift = ctx->electronShift * shifts[grid] * PetscSqrtReal(data[0].kT_m) / ctx->v_0; // shifts to not matter!!!!
    115:   }
    116:   PetscCall(DMProjectFunction(dm, time, initu, (void **)mctxs, INSERT_ALL_VALUES, X));
    117:   PetscFunctionReturn(PETSC_SUCCESS);
    118: }

    120: typedef enum {
    121:   E_PAR_IDX,
    122:   E_PERP_IDX,
    123:   I_PAR_IDX,
    124:   I_PERP_IDX,
    125:   NUM_TEMPS
    126: } TemperatureIDX;

    128: /* --------------------  Evaluate NRL Function F(x) (analytical solutions exist for this) --------------------- */
    129: static PetscReal n_cm3[2] = {0, 0};
    130: PetscErrorCode   FormFunction(TS ts, PetscReal tdummy, Vec X, Vec F, void *ptr)
    131: {
    132:   LandauCtx         *ctx = (LandauCtx *)ptr; /* user-defined application context */
    133:   PetscScalar       *f;
    134:   const PetscScalar *x;
    135:   const PetscReal    k_B = 1.6e-12, e_cgs = 4.8e-10, proton_mass = 9.1094e-28, m_cgs[2] = {proton_mass, proton_mass * ctx->masses[1] / ctx->masses[0]}; // erg/eV, e, m as per NRL;
    136:   PetscReal          AA, v_bar_ab, vTe, t1, TeDiff, Te, Ti, Tdiff;

    138:   PetscFunctionBeginUser;
    139:   PetscCall(VecGetArrayRead(X, &x));
    140:   Te = PetscRealPart(2 * x[E_PERP_IDX] + x[E_PAR_IDX]) / 3, Ti = PetscRealPart(2 * x[I_PERP_IDX] + x[I_PAR_IDX]) / 3;
    141:   // thermalization from NRL Plasma formulary, assume Z = 1, mu = 2, n_i = n_e
    142:   v_bar_ab = 1.8e-19 * PetscSqrtReal(m_cgs[0] * m_cgs[1]) * n_cm3[0] * ctx->lambdas[0][1] * PetscPowReal(m_cgs[0] * Ti + m_cgs[1] * Te, -1.5);
    143:   PetscCall(VecGetArray(F, &f));
    144:   for (PetscInt ii = 0; ii < 2; ii++) {
    145:     PetscReal tPerp = PetscRealPart(x[2 * ii + E_PERP_IDX]), tPar = PetscRealPart(x[2 * ii + E_PAR_IDX]), ff;
    146:     TeDiff = tPerp - tPar;
    147:     AA     = tPerp / tPar - 1;
    148:     if (AA < 0) ff = PetscAtanhReal(PetscSqrtReal(-AA)) / PetscSqrtReal(-AA);
    149:     else ff = PetscAtanReal(PetscSqrtReal(AA)) / PetscSqrtReal(AA);
    150:     t1 = (-3 + (AA + 3) * ff) / PetscSqr(AA);
    151:     //PetscReal vTeB = 8.2e-7 * n_cm3[0] * ctx->lambdas[0][1] * PetscPowReal(Te, -1.5);
    152:     vTe = 2 * PetscSqrtReal(PETSC_PI / m_cgs[ii]) * PetscSqr(PetscSqr(e_cgs)) * n_cm3[0] * ctx->lambdas[0][1] * PetscPowReal(PetscRealPart(k_B * x[E_PAR_IDX]), -1.5) * t1;
    153:     t1  = vTe * TeDiff; // * 2; // scaling from NRL that makes it fit pretty good

    155:     f[2 * ii + E_PAR_IDX]  = 2 * t1; // par
    156:     f[2 * ii + E_PERP_IDX] = -t1;    // perp
    157:     Tdiff                  = (ii == 0) ? (Ti - Te) : (Te - Ti);
    158:     f[2 * ii + E_PAR_IDX] += v_bar_ab * Tdiff;
    159:     f[2 * ii + E_PERP_IDX] += v_bar_ab * Tdiff;
    160:   }
    161:   PetscCall(VecRestoreArrayRead(X, &x));
    162:   PetscCall(VecRestoreArray(F, &f));
    163:   PetscFunctionReturn(PETSC_SUCCESS);
    164: }

    166: /* --------------------  Form initial approximation ----------------- */
    167: static PetscReal T0[4] = {300, 390, 200, 260};
    168: PetscErrorCode   createVec_NRL(LandauCtx *ctx, Vec *vec)
    169: {
    170:   PetscScalar *x;
    171:   Vec          Temps;

    173:   PetscFunctionBeginUser;
    174:   PetscCall(VecCreateSeq(PETSC_COMM_SELF, NUM_TEMPS, &Temps));
    175:   PetscCall(VecGetArray(Temps, &x));
    176:   for (PetscInt i = 0; i < NUM_TEMPS; i++) x[i] = T0[i];
    177:   PetscCall(VecRestoreArray(Temps, &x));
    178:   *vec = Temps;
    179:   PetscFunctionReturn(PETSC_SUCCESS);
    180: }

    182: PetscErrorCode createTS_NRL(LandauCtx *ctx, Vec Temps)
    183: {
    184:   TSAdapt adapt;
    185:   TS      ts;

    187:   PetscFunctionBeginUser;
    188:   PetscCall(TSCreate(PETSC_COMM_SELF, &ts));
    189:   ctx->data = (void *)ts; // 'data' is for applications (eg, monitors)
    190:   PetscCall(TSSetApplicationContext(ts, ctx));
    191:   PetscCall(TSSetType(ts, TSRK));
    192:   PetscCall(TSSetRHSFunction(ts, NULL, FormFunction, ctx));
    193:   PetscCall(TSSetSolution(ts, Temps));
    194:   PetscCall(TSRKSetType(ts, TSRK2A));
    195:   PetscCall(TSSetOptionsPrefix(ts, "nrl_"));
    196:   PetscCall(TSSetFromOptions(ts));
    197:   PetscCall(TSGetAdapt(ts, &adapt));
    198:   PetscCall(TSAdaptSetType(adapt, TSADAPTNONE));
    199:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
    200:   PetscCall(TSSetStepNumber(ts, 0));
    201:   PetscCall(TSSetMaxSteps(ts, 1));
    202:   PetscCall(TSSetTime(ts, 0));
    203:   PetscFunctionReturn(PETSC_SUCCESS);
    204: }

    206: PetscErrorCode Monitor_nrl(TS ts, PetscInt stepi, PetscReal time, Vec X, void *actx)
    207: {
    208:   const PetscScalar *x;
    209:   LandauCtx         *ctx = (LandauCtx *)actx; /* user-defined application context */

    211:   PetscFunctionBeginUser;
    212:   if (stepi % 100 == 0) {
    213:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "nrl-step %d time= %g ", (int)stepi, (double)(time / ctx->t_0)));
    214:     PetscCall(VecGetArrayRead(X, &x));
    215:     for (PetscInt i = 0; i < NUM_TEMPS; i++) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%g ", (double)PetscRealPart(x[i])));
    216:     PetscCall(VecRestoreArrayRead(X, &x));
    217:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n"));
    218:   }
    219:   PetscFunctionReturn(PETSC_SUCCESS);
    220: }

    222: PetscErrorCode Monitor(TS ts, PetscInt stepi, PetscReal time, Vec X, void *actx)
    223: {
    224:   LandauCtx *ctx      = (LandauCtx *)actx; /* user-defined application context */
    225:   TS         ts_nrl   = (TS)ctx->data;
    226:   PetscInt   printing = 0, logT;

    228:   PetscFunctionBeginUser;
    229:   if (ctx->verbose > 0) { // hacks to generate sparse data (eg, use '-dm_landau_verbose 1' and '-dm_landau_verbose -1' to get all steps printed)
    230:     PetscReal dt;
    231:     PetscCall(TSGetTimeStep(ts, &dt));
    232:     logT = (PetscInt)PetscLog2Real(time / dt);
    233:     if (logT < 0) logT = 0;
    234:     ctx->verbose = PetscPowInt(2, logT) / 2;
    235:     if (ctx->verbose == 0) ctx->verbose = 1;
    236:   }
    237:   if (ctx->verbose) {
    238:     TSConvergedReason reason;
    239:     PetscCall(TSGetConvergedReason(ts, &reason));
    240:     if (stepi % ctx->verbose == 0 || reason || stepi == 1 || ctx->verbose < 0) {
    241:       PetscInt nDMs, id;
    242:       DM       pack;
    243:       Vec     *XsubArray = NULL;
    244:       printing           = 1;
    245:       PetscCall(TSGetDM(ts, &pack));
    246:       PetscCall(DMCompositeGetNumberDM(pack, &nDMs));
    247:       PetscCall(DMGetOutputSequenceNumber(ctx->plex[0], &id, NULL));
    248:       PetscCall(DMSetOutputSequenceNumber(ctx->plex[0], id + 1, time));
    249:       PetscCall(DMSetOutputSequenceNumber(ctx->plex[1], id + 1, time));
    250:       PetscCall(PetscInfo(pack, "ex1 plot step %" PetscInt_FMT ", time = %g\n", id, (double)time));
    251:       PetscCall(PetscMalloc(sizeof(*XsubArray) * nDMs, &XsubArray));
    252:       PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, XsubArray)); // read only
    253:       PetscCall(VecViewFromOptions(XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], NULL, "-ex1_vec_view_e"));
    254:       PetscCall(VecViewFromOptions(XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 1)], NULL, "-ex1_vec_view_i"));
    255:       // temps
    256:       for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
    257:         PetscDS     prob;
    258:         DM          dm      = ctx->plex[grid];
    259:         PetscScalar user[2] = {0, 0}, tt[1];
    260:         PetscReal   vz_0 = 0, n, energy, e_perp, e_par, m_s = ctx->masses[ctx->species_offset[grid]];
    261:         Vec         Xloc = XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, grid)];
    262:         PetscCall(DMGetDS(dm, &prob));
    263:         /* get n */
    264:         PetscCall(PetscDSSetObjective(prob, 0, &f0_n));
    265:         PetscCall(DMPlexComputeIntegralFEM(dm, Xloc, tt, NULL));
    266:         n = PetscRealPart(tt[0]);
    267:         /* get vz */
    268:         PetscCall(PetscDSSetObjective(prob, 0, &f0_vz));
    269:         PetscCall(DMPlexComputeIntegralFEM(dm, Xloc, tt, NULL));
    270:         user[0] = vz_0 = PetscRealPart(tt[0]) / n;
    271:         /* energy temp */
    272:         PetscCall(PetscDSSetConstants(prob, 2, user));
    273:         PetscCall(PetscDSSetObjective(prob, 0, &f0_v2_shift));
    274:         PetscCall(DMPlexComputeIntegralFEM(dm, Xloc, tt, ctx));
    275:         energy = PetscRealPart(tt[0]) * ctx->v_0 * ctx->v_0 * m_s / n / 3; // scale?
    276:         energy *= kev_joul * 1000;                                         // T eV
    277:         /* energy temp - par */
    278:         PetscCall(PetscDSSetConstants(prob, 2, user));
    279:         PetscCall(PetscDSSetObjective(prob, 0, &f0_v2_par_shift));
    280:         PetscCall(DMPlexComputeIntegralFEM(dm, Xloc, tt, ctx));
    281:         e_par = PetscRealPart(tt[0]) * ctx->v_0 * ctx->v_0 * m_s / n;
    282:         e_par *= kev_joul * 1000; // eV
    283:         /* energy temp - perp */
    284:         PetscCall(PetscDSSetConstants(prob, 2, user));
    285:         PetscCall(PetscDSSetObjective(prob, 0, &f0_v2_perp));
    286:         PetscCall(DMPlexComputeIntegralFEM(dm, Xloc, tt, ctx));
    287:         e_perp = PetscRealPart(tt[0]) * ctx->v_0 * ctx->v_0 * m_s / n / 2;
    288:         e_perp *= kev_joul * 1000; // eV
    289:         if (grid == 0) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step %4d) time= %e temperature (eV): ", (int)stepi, (double)time));
    290:         PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%s T= %9.4e T_par= %9.4e T_perp= %9.4e ", (grid == 0) ? "electron:" : ";ion:", (double)energy, (double)e_par, (double)e_perp));
    291:         if (n_cm3[grid] == 0) n_cm3[grid] = ctx->n_0 * n * 1e-6; // does not change m^3 --> cm^3
    292:       }
    293:       // cleanup
    294:       PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, XsubArray));
    295:       PetscCall(PetscFree(XsubArray));
    296:     }
    297:   }
    298:   /* evolve NRL data, end line */
    299:   if (n_cm3[NUM_TEMPS / 2 - 1] < 0 && ts_nrl) {
    300:     PetscCall(TSDestroy(&ts_nrl));
    301:     ctx->data = NULL;
    302:     if (printing) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nSTOP printing NRL Ts\n"));
    303:   } else if (ts_nrl) {
    304:     const PetscScalar *x;
    305:     PetscReal          dt_real, dt;
    306:     Vec                U;
    307:     PetscCall(TSGetTimeStep(ts, &dt)); // dt for NEXT time step
    308:     dt_real = dt * ctx->t_0;
    309:     PetscCall(TSGetSolution(ts_nrl, &U));
    310:     if (printing) {
    311:       PetscCall(VecGetArrayRead(U, &x));
    312:       PetscCall(PetscPrintf(PETSC_COMM_WORLD, "NRL_e_par= %9.4e NRL_e_perp= %9.4e ", (double)PetscRealPart(x[E_PAR_IDX]), (double)PetscRealPart(x[E_PERP_IDX])));
    313:       PetscCall(PetscPrintf(PETSC_COMM_WORLD, "NRL_i_par= %9.4e NRL_i_perp= %9.4e\n", (double)PetscRealPart(x[I_PAR_IDX]), (double)PetscRealPart(x[I_PERP_IDX])));
    314:       /* if (n_cm3[0] > 0) { */
    315:       /* } else PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n")); */
    316:       PetscCall(VecRestoreArrayRead(U, &x));
    317:     }
    318:     // we have the next time step, so need to advance now
    319:     PetscCall(TSSetTimeStep(ts_nrl, dt_real));
    320:     PetscCall(TSSetMaxSteps(ts_nrl, stepi + 1)); // next step
    321:     PetscCall(TSSolve(ts_nrl, NULL));
    322:   } else if (printing) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n"));
    323:   if (printing) PetscCall(DMPlexLandauPrintNorms(X, stepi /*id + 1*/));
    324:   PetscFunctionReturn(PETSC_SUCCESS);
    325: }

    327: int main(int argc, char **argv)
    328: {
    329:   DM          pack;
    330:   Vec         X;
    331:   PetscInt    dim = 2, nDMs;
    332:   TS          ts, ts_nrl = NULL;
    333:   Mat         J;
    334:   Vec        *XsubArray = NULL;
    335:   LandauCtx  *ctx;
    336:   PetscMPIInt rank;
    337:   PetscBool   use_nrl   = PETSC_TRUE;
    338:   PetscBool   print_nrl = PETSC_FALSE;
    339:   PetscReal   dt0;

    341:   PetscFunctionBeginUser;
    342:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
    343:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
    344:   if (rank) { /* turn off output stuff for duplicate runs */
    345:     PetscCall(PetscOptionsClearValue(NULL, "-ex1_dm_view_e"));
    346:     PetscCall(PetscOptionsClearValue(NULL, "-ex1_dm_view_i"));
    347:     PetscCall(PetscOptionsClearValue(NULL, "-ex1_vec_view_e"));
    348:     PetscCall(PetscOptionsClearValue(NULL, "-ex1_vec_view_i"));
    349:     PetscCall(PetscOptionsClearValue(NULL, "-info"));
    350:     PetscCall(PetscOptionsClearValue(NULL, "-snes_converged_reason"));
    351:     PetscCall(PetscOptionsClearValue(NULL, "-pc_bjkokkos_ksp_converged_reason"));
    352:     PetscCall(PetscOptionsClearValue(NULL, "-ksp_converged_reason"));
    353:     PetscCall(PetscOptionsClearValue(NULL, "-ts_adapt_monitor"));
    354:     PetscCall(PetscOptionsClearValue(NULL, "-ts_monitor"));
    355:     PetscCall(PetscOptionsClearValue(NULL, "-snes_monitor"));
    356:   }
    357:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-dim", &dim, NULL));
    358:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-use_nrl", &use_nrl, NULL));
    359:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-print_nrl", &print_nrl, NULL));
    360:   /* Create a mesh */
    361:   PetscCall(DMPlexLandauCreateVelocitySpace(PETSC_COMM_SELF, dim, "", &X, &J, &pack));
    362:   PetscCall(DMSetUp(pack));
    363:   PetscCall(DMGetApplicationContext(pack, &ctx));
    364:   PetscCheck(ctx->num_grids == 2, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "Must have two grids: use '-dm_landau_num_species_grid 1,1'");
    365:   PetscCheck(ctx->num_species == 2, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "Must have two species: use '-dm_landau_num_species_grid 1,1'");
    366:   PetscCall(DMCompositeGetNumberDM(pack, &nDMs));
    367:   /* output plot names */
    368:   PetscCall(PetscMalloc(sizeof(*XsubArray) * nDMs, &XsubArray));
    369:   PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, XsubArray)); // read only
    370:   PetscCall(PetscObjectSetName((PetscObject)XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], 0 == 0 ? "ue" : "ui"));
    371:   PetscCall(PetscObjectSetName((PetscObject)XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 1)], 1 == 0 ? "ue" : "ui"));
    372:   /* add bimaxwellian anisotropic test */
    373:   for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) {
    374:     for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
    375:       PetscReal shifts[2];
    376:       PetscCall(SetMaxwellians(ctx->plex[grid], XsubArray[LAND_PACK_IDX(b_id, grid)], 0.0, ctx->thermal_temps, ctx->n, grid, shifts, ctx));
    377:     }
    378:   }
    379:   PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, XsubArray));
    380:   PetscCall(PetscFree(XsubArray));
    381:   /* plot */
    382:   PetscCall(DMSetOutputSequenceNumber(ctx->plex[0], -1, 0.0));
    383:   PetscCall(DMSetOutputSequenceNumber(ctx->plex[1], -1, 0.0));
    384:   PetscCall(DMViewFromOptions(ctx->plex[0], NULL, "-ex1_dm_view_e"));
    385:   PetscCall(DMViewFromOptions(ctx->plex[1], NULL, "-ex1_dm_view_i"));
    386:   /* Create timestepping solver context */
    387:   PetscCall(TSCreate(PETSC_COMM_SELF, &ts));
    388:   PetscCall(TSSetDM(ts, pack));
    389:   PetscCall(TSSetIFunction(ts, NULL, DMPlexLandauIFunction, NULL));
    390:   PetscCall(TSSetIJacobian(ts, J, J, DMPlexLandauIJacobian, NULL));
    391:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
    392:   PetscCall(TSSetFromOptions(ts));
    393:   PetscCall(TSSetSolution(ts, X));
    394:   PetscCall(TSMonitorSet(ts, Monitor, ctx, NULL));
    395:   /* Create NRL timestepping */
    396:   if (use_nrl || print_nrl) {
    397:     Vec NRL_vec;
    398:     PetscCall(createVec_NRL(ctx, &NRL_vec));
    399:     PetscCall(createTS_NRL(ctx, NRL_vec));
    400:     PetscCall(VecDestroy(&NRL_vec));
    401:   } else ctx->data = NULL;
    402:   /* solve */
    403:   PetscCall(TSGetTimeStep(ts, &dt0));
    404:   PetscCall(TSSetTime(ts, dt0 / 2));
    405:   PetscCall(TSSolve(ts, X));
    406:   /* test add field method & output */
    407:   PetscCall(DMPlexLandauAccess(pack, X, landau_field_print_access_callback, NULL));
    408:   // run NRL in separate TS
    409:   ts_nrl = (TS)ctx->data;
    410:   if (print_nrl) {
    411:     PetscReal    finalTime, dt_real, tstart = dt0 * ctx->t_0 / 2; // hack
    412:     Vec          U;
    413:     PetscScalar *x;
    414:     PetscInt     nsteps;
    415:     dt_real = dt0 * ctx->t_0;
    416:     PetscCall(TSSetTimeStep(ts_nrl, dt_real));
    417:     PetscCall(TSGetTime(ts, &finalTime));
    418:     finalTime *= ctx->t_0;
    419:     PetscCall(TSSetMaxTime(ts_nrl, finalTime));
    420:     nsteps = (PetscInt)(finalTime / dt_real) + 1;
    421:     PetscCall(TSSetMaxSteps(ts_nrl, nsteps));
    422:     PetscCall(TSSetStepNumber(ts_nrl, 0));
    423:     PetscCall(TSSetTime(ts_nrl, tstart));
    424:     PetscCall(TSGetSolution(ts_nrl, &U));
    425:     PetscCall(VecGetArray(U, &x));
    426:     for (PetscInt i = 0; i < NUM_TEMPS; i++) x[i] = T0[i];
    427:     PetscCall(VecRestoreArray(U, &x));
    428:     PetscCall(TSMonitorSet(ts_nrl, Monitor_nrl, ctx, NULL));
    429:     PetscCall(TSSolve(ts_nrl, NULL));
    430:   }
    431:   /* clean up */
    432:   PetscCall(TSDestroy(&ts));
    433:   PetscCall(TSDestroy(&ts_nrl));
    434:   PetscCall(VecDestroy(&X));
    435:   PetscCall(DMPlexLandauDestroyVelocitySpace(&pack));
    436:   PetscCall(PetscFinalize());
    437:   return 0;
    438: }

    440: /*TEST
    441:   testset:
    442:     requires: p4est !complex double defined(PETSC_USE_DMLANDAU_2D)
    443:     output_file: output/ex1_0.out
    444:     filter: grep -v "DM"
    445:     args: -dm_landau_amr_levels_max 0,2 -dm_landau_amr_post_refine 0 -dm_landau_amr_re_levels 2 -dm_landau_domain_radius 6,6 -dm_landau_electron_shift 1.5 -dm_landau_ion_charges 1 -dm_landau_ion_masses 2 -dm_landau_n 1,1 -dm_landau_n_0 1e20 -dm_landau_num_cells 2,4 -dm_landau_num_species_grid 1,1 -dm_landau_re_radius 2 -use_nrl true -print_nrl false -dm_landau_thermal_temps .3,.2 -dm_landau_type p4est -dm_landau_verbose -1 -dm_preallocate_only false -ex1_dm_view_e -ksp_type preonly -pc_type lu -petscspace_degree 3 -snes_converged_reason -snes_rtol 1.e-14 -snes_stol 1.e-14 -ts_adapt_clip .5,1.5 -ts_adapt_dt_max 5 -ts_adapt_monitor -ts_adapt_scale_solve_failed 0.5 -ts_arkimex_type 1bee -ts_time_step .01 -ts_max_snes_failures unlimited -ts_max_steps 1 -ts_max_time 8 -ts_monitor -ts_rtol 1e-2 -ts_type arkimex
    446:     test:
    447:       suffix: cpu
    448:       args: -dm_landau_device_type cpu -dm_landau_use_relativistic_corrections
    449:     test:
    450:       suffix: kokkos
    451:       requires: kokkos_kernels !defined(PETSC_HAVE_CUDA_CLANG)
    452:       args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos

    454:   testset:
    455:     requires: !complex defined(PETSC_USE_DMLANDAU_2D) p4est
    456:     args: -dm_landau_type p4est -dm_landau_num_species_grid 1,1 -dm_landau_n 1,1 -dm_landau_thermal_temps 2,1 -dm_landau_ion_charges 1 -dm_landau_ion_masses 2 -petscspace_degree 2 -ts_type beuler -ts_time_step .1 -ts_max_steps 1 -dm_landau_verbose 2 -ksp_type preonly -pc_type lu -dm_landau_device_type cpu -use_nrl false -print_nrl -snes_rtol 1.e-14 -snes_stol 1.e-14 -snes_converged_reason -dm_landau_device_type cpu
    457:     nsize: 1
    458:     test:
    459:       suffix: sphere
    460:       args: -dm_landau_sphere -dm_landau_amr_levels_max 1,1 -dm_landau_sphere_inner_radius_90degree_scale .55 -dm_landau_sphere_inner_radius_45degree_scale .5
    461:     test:
    462:       suffix: re
    463:       args: -dm_landau_num_cells 4,4 -dm_landau_amr_levels_max 0,2 -dm_landau_z_radius_pre 2.5 -dm_landau_z_radius_post 3.75 -dm_landau_amr_z_refine_pre 1 -dm_landau_amr_z_refine_post 1 -dm_landau_electron_shift 1.25 -info :vec

    465: TEST*/