xref: /petsc/src/ts/utils/dmplexlandau/tutorials/ex2.c (revision 48a46eb9bd028bec07ec0f396b1a3abb43f14558)
1e0eea495SMark static char help[] = "Runaway electron model with Landau collision operator\n\n";
2e0eea495SMark 
3e0eea495SMark #include <petscdmplex.h>
4e0eea495SMark #include <petsclandau.h>
5e0eea495SMark #include <petscts.h>
6e0eea495SMark #include <petscds.h>
78a6f2e61SMark Adams #include <petscdmcomposite.h>
8c3e4dd79SMark Adams #include "petsc/private/petscimpl.h"
9e0eea495SMark 
10419c2fb8Smarkadams4 #if defined(PETSC_HAVE_CUDA_NVTX)
11419c2fb8Smarkadams4 #include <nvToolsExt.h>
12419c2fb8Smarkadams4 #endif
13419c2fb8Smarkadams4 
14e0eea495SMark /* data for runaway electron model */
15e0eea495SMark typedef struct REctx_struct {
168a6f2e61SMark Adams   PetscErrorCode (*test)(TS, Vec, PetscInt, PetscReal, PetscBool, LandauCtx *, struct REctx_struct *);
17e0eea495SMark   PetscErrorCode (*impuritySrcRate)(PetscReal, PetscReal *, LandauCtx *);
18e0eea495SMark   PetscErrorCode (*E)(Vec, Vec, PetscInt, PetscReal, LandauCtx *, PetscReal *);
19e0eea495SMark   PetscReal T_cold;        /* temperature of newly ionized electrons and impurity ions */
20e0eea495SMark   PetscReal ion_potential; /* ionization potential of impurity */
21e0eea495SMark   PetscReal Ne_ion;        /* effective number of electrons shed in ioization of impurity */
22e0eea495SMark   PetscReal Ez_initial;
23e0eea495SMark   PetscReal L; /* inductance */
24e0eea495SMark   Vec       X_0;
25e0eea495SMark   PetscInt  imp_idx; /* index for impurity ionizing sink */
26e0eea495SMark   PetscReal pulse_start;
27e0eea495SMark   PetscReal pulse_width;
28e0eea495SMark   PetscReal pulse_rate;
29e0eea495SMark   PetscReal current_rate;
30e0eea495SMark   PetscInt  plotIdx;
31e0eea495SMark   PetscInt  plotStep;
32e0eea495SMark   PetscInt  idx; /* cache */
33e0eea495SMark   PetscReal j;   /* cache */
34e0eea495SMark   PetscReal plotDt;
35e0eea495SMark   PetscBool plotting;
36e0eea495SMark   PetscBool use_spitzer_eta;
378a6f2e61SMark Adams   PetscInt  print_period;
388fdabdddSMark Adams   PetscInt  grid_view_idx;
39e0eea495SMark } REctx;
40e0eea495SMark 
41e0eea495SMark static const PetscReal kev_joul = 6.241506479963235e+15; /* 1/1000e */
42e0eea495SMark 
4390d163c9SMark Adams #define RE_CUT 3.
44e0eea495SMark /* < v, u_re * v * q > */
459371c9d4SSatish Balay 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) {
46e0eea495SMark   PetscReal n_e = PetscRealPart(u[0]);
47e0eea495SMark   if (dim == 2) {
48e0eea495SMark     if (x[1] > RE_CUT || x[1] < -RE_CUT) {                    /* simply a cutoff for REs. v_|| > 3 v(T_e) */
49e0eea495SMark       *f0 = n_e * 2. * PETSC_PI * x[0] * x[1] * constants[0]; /* n * r * v_|| * q */
50e0eea495SMark     } else {
51e0eea495SMark       *f0 = 0;
52e0eea495SMark     }
53e0eea495SMark   } else {
54e0eea495SMark     if (x[2] > RE_CUT || x[2] < -RE_CUT) { /* simply a cutoff for REs. v_|| > 3 v(T_e) */
55e0eea495SMark       *f0 = n_e * x[2] * constants[0];
56e0eea495SMark     } else {
57e0eea495SMark       *f0 = 0;
58e0eea495SMark     }
59e0eea495SMark   }
60e0eea495SMark }
61e0eea495SMark 
62e0eea495SMark /* sum < v, u*v*q > */
639371c9d4SSatish Balay 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) {
64e0eea495SMark   PetscInt ii;
65e0eea495SMark   f0[0] = 0;
66e0eea495SMark   if (dim == 2) {
678a6f2e61SMark Adams     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 */
68e0eea495SMark   } else {
698a6f2e61SMark Adams     for (ii = 0; ii < Nf; ii++) f0[0] += u[ii] * x[2] * q[ii]; /* n * v_|| * q  * v_0 */
70e0eea495SMark   }
71e0eea495SMark }
72e0eea495SMark 
73e0eea495SMark /* < v, n_e > */
749371c9d4SSatish Balay 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) {
75e0eea495SMark   PetscInt ii = (PetscInt)PetscRealPart(constants[0]);
76e0eea495SMark   if (dim == 2) f0[0] = 2. * PETSC_PI * x[0] * u[ii];
779371c9d4SSatish Balay   else { f0[0] = u[ii]; }
78e0eea495SMark }
79e0eea495SMark 
80e0eea495SMark /* < v, n_e v_|| > */
819371c9d4SSatish Balay 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) {
82e0eea495SMark   PetscInt ii = (PetscInt)PetscRealPart(constants[0]);
83e0eea495SMark   if (dim == 2) f0[0] = u[ii] * 2. * PETSC_PI * x[0] * x[1]; /* n r v_|| */
849371c9d4SSatish Balay   else { f0[0] = u[ii] * x[2]; /* n v_|| */ }
85e0eea495SMark }
86e0eea495SMark 
87e0eea495SMark /* < v, n_e (v-shift) > */
889371c9d4SSatish Balay 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) {
898a6f2e61SMark Adams   PetscReal vz = numConstants > 0 ? PetscRealPart(constants[0]) : 0;
90e0eea495SMark   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 */
919371c9d4SSatish Balay   else { *f0 = u[0] * PetscSqrtReal(x[0] * x[0] + x[1] * x[1] + (x[2] - vz) * (x[2] - vz)); /* n v */ }
92e0eea495SMark }
93e0eea495SMark 
94e0eea495SMark /* CalculateE - Calculate the electric field  */
95e0eea495SMark /*  T        -- Electron temperature  */
96e0eea495SMark /*  n        -- Electron density  */
97e0eea495SMark /*  lnLambda --   */
98e0eea495SMark /*  eps0     --  */
99e0eea495SMark /*  E        -- output E, input \hat E */
1009371c9d4SSatish Balay static PetscReal CalculateE(PetscReal Tev, PetscReal n, PetscReal lnLambda, PetscReal eps0, PetscReal *E) {
101e0eea495SMark   PetscReal c, e, m;
102e0eea495SMark 
103e0eea495SMark   PetscFunctionBegin;
104cefb98e8SMark Adams   c = 299792458.0;
105e0eea495SMark   e = 1.602176e-19;
106e0eea495SMark   m = 9.10938e-31;
107e0eea495SMark   if (1) {
108e0eea495SMark     double Ec, Ehat = *E, betath = PetscSqrtReal(2 * Tev * e / (m * c * c)), j0 = Ehat * 7 / (PetscSqrtReal(2) * 2) * PetscPowReal(betath, 3) * n * e * c;
109e0eea495SMark     Ec = n * lnLambda * PetscPowReal(e, 3) / (4 * PETSC_PI * PetscPowReal(eps0, 2) * m * c * c);
110e0eea495SMark     *E = Ec;
111e0eea495SMark     PetscPrintf(PETSC_COMM_WORLD, "CalculateE j0=%g Ec = %g\n", j0, Ec);
112e0eea495SMark   } else {
113e0eea495SMark     PetscReal Ed, vth;
114e0eea495SMark     vth = PetscSqrtReal(8 * Tev * e / (m * PETSC_PI));
115e0eea495SMark     Ed  = n * lnLambda * PetscPowReal(e, 3) / (4 * PETSC_PI * PetscPowReal(eps0, 2) * m * vth * vth);
116e0eea495SMark     *E  = Ed;
117e0eea495SMark   }
118e0eea495SMark   PetscFunctionReturn(0);
119e0eea495SMark }
120e0eea495SMark 
1219371c9d4SSatish Balay static PetscReal Spitzer(PetscReal m_e, PetscReal e, PetscReal Z, PetscReal epsilon0, PetscReal lnLam, PetscReal kTe_joules) {
122e0eea495SMark   PetscReal Fz = (1 + 1.198 * Z + 0.222 * Z * Z) / (1 + 2.966 * Z + 0.753 * Z * Z), eta;
123e0eea495SMark   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);
124e0eea495SMark   return eta;
125e0eea495SMark }
126e0eea495SMark 
127e0eea495SMark /*  */
1289371c9d4SSatish Balay static PetscErrorCode testNone(TS ts, Vec X, PetscInt stepi, PetscReal time, PetscBool islast, LandauCtx *ctx, REctx *rectx) {
129e0eea495SMark   PetscFunctionBeginUser;
130e0eea495SMark   PetscFunctionReturn(0);
131e0eea495SMark }
132e0eea495SMark 
133e0eea495SMark /*  */
1349371c9d4SSatish Balay static PetscErrorCode testSpitzer(TS ts, Vec X, PetscInt stepi, PetscReal time, PetscBool islast, LandauCtx *ctx, REctx *rectx) {
135f37ccb5fSMark Adams   PetscInt          ii, nDMs;
136e0eea495SMark   PetscDS           prob;
137a587d139SMark   static PetscReal  old_ratio = 1e10;
138a587d139SMark   TSConvergedReason reason;
139a587d139SMark   PetscReal         J, J_re, spit_eta, Te_kev = 0, E, ratio, Z, n_e, v, v2;
1408a6f2e61SMark Adams   PetscScalar       user[2] = {0., ctx->charges[0]}, q[LANDAU_MAX_SPECIES], tt[LANDAU_MAX_SPECIES], vz;
141930e68a5SMark Adams   PetscReal         dt;
1428a6f2e61SMark Adams   DM                pack, plexe = ctx->plex[0], plexi = (ctx->num_grids == 1) ? NULL : ctx->plex[1];
143f37ccb5fSMark Adams   Vec              *XsubArray;
144e0eea495SMark 
145e0eea495SMark   PetscFunctionBeginUser;
14663a3b9bcSJacob Faibussowitsch   PetscCheck(ctx->num_species == 2, PETSC_COMM_SELF, PETSC_ERR_PLIB, "ctx->num_species %" PetscInt_FMT " != 2", ctx->num_species);
1479566063dSJacob Faibussowitsch   PetscCall(VecGetDM(X, &pack));
1483c633725SBarry Smith   PetscCheck(pack, PETSC_COMM_SELF, PETSC_ERR_PLIB, "no DM");
1499566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetNumberDM(pack, &nDMs));
15063a3b9bcSJacob Faibussowitsch   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);
1519566063dSJacob Faibussowitsch   PetscCall(PetscMalloc(sizeof(*XsubArray) * nDMs, &XsubArray));
1529566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, XsubArray)); // read only
1539566063dSJacob Faibussowitsch   PetscCall(TSGetTimeStep(ts, &dt));
1548a6f2e61SMark Adams   /* get current for each grid */
1558a6f2e61SMark Adams   for (ii = 0; ii < ctx->num_species; ii++) q[ii] = ctx->charges[ii];
1569566063dSJacob Faibussowitsch   PetscCall(DMGetDS(plexe, &prob));
1579566063dSJacob Faibussowitsch   PetscCall(PetscDSSetConstants(prob, 2, &q[0]));
1589566063dSJacob Faibussowitsch   PetscCall(PetscDSSetObjective(prob, 0, &f0_jz_sum));
1599566063dSJacob Faibussowitsch   PetscCall(DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL));
160e0eea495SMark   J = -ctx->n_0 * ctx->v_0 * PetscRealPart(tt[0]);
1618fdabdddSMark Adams   if (plexi) { // add first (only) ion
1629566063dSJacob Faibussowitsch     PetscCall(DMGetDS(plexi, &prob));
1639566063dSJacob Faibussowitsch     PetscCall(PetscDSSetConstants(prob, 1, &q[1]));
1649566063dSJacob Faibussowitsch     PetscCall(PetscDSSetObjective(prob, 0, &f0_jz_sum));
1659566063dSJacob Faibussowitsch     PetscCall(DMPlexComputeIntegralFEM(plexi, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 1)], tt, NULL));
1668a6f2e61SMark Adams     J += -ctx->n_0 * ctx->v_0 * PetscRealPart(tt[0]);
1678a6f2e61SMark Adams   }
168a587d139SMark   /* get N_e */
1699566063dSJacob Faibussowitsch   PetscCall(DMGetDS(plexe, &prob));
1709566063dSJacob Faibussowitsch   PetscCall(PetscDSSetConstants(prob, 1, user));
1719566063dSJacob Faibussowitsch   PetscCall(PetscDSSetObjective(prob, 0, &f0_n));
1729566063dSJacob Faibussowitsch   PetscCall(DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL));
173a587d139SMark   n_e = PetscRealPart(tt[0]) * ctx->n_0;
174a587d139SMark   /* Z */
175a587d139SMark   Z   = -ctx->charges[1] / ctx->charges[0];
176a587d139SMark   /* remove drift */
1778a6f2e61SMark Adams   if (0) {
1788a6f2e61SMark Adams     user[0] = 0; // electrons
1799566063dSJacob Faibussowitsch     PetscCall(DMGetDS(plexe, &prob));
1809566063dSJacob Faibussowitsch     PetscCall(PetscDSSetConstants(prob, 1, user));
1819566063dSJacob Faibussowitsch     PetscCall(PetscDSSetObjective(prob, 0, &f0_vz));
1829566063dSJacob Faibussowitsch     PetscCall(DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL));
183a587d139SMark     vz = ctx->n_0 * PetscRealPart(tt[0]) / n_e; /* non-dimensional */
184a587d139SMark   } else vz = 0;
185a587d139SMark   /* thermal velocity */
1869566063dSJacob Faibussowitsch   PetscCall(DMGetDS(plexe, &prob));
1879566063dSJacob Faibussowitsch   PetscCall(PetscDSSetConstants(prob, 1, &vz));
1889566063dSJacob Faibussowitsch   PetscCall(PetscDSSetObjective(prob, 0, &f0_ve_shift));
1899566063dSJacob Faibussowitsch   PetscCall(DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL));
190a587d139SMark   v        = ctx->n_0 * ctx->v_0 * PetscRealPart(tt[0]) / n_e;                                           /* remove number density to get velocity */
191a587d139SMark   v2       = PetscSqr(v);                                                                                /* use real space: m^2 / s^2 */
192a587d139SMark   Te_kev   = (v2 * ctx->masses[0] * PETSC_PI / 8) * kev_joul;                                            /* temperature in kev */
193a587d139SMark   spit_eta = Spitzer(ctx->masses[0], -ctx->charges[0], Z, ctx->epsilon0, ctx->lnLam, Te_kev / kev_joul); /* kev --> J (kT) */
1948a6f2e61SMark Adams   if (0) {
1959566063dSJacob Faibussowitsch     PetscCall(DMGetDS(plexe, &prob));
1969566063dSJacob Faibussowitsch     PetscCall(PetscDSSetConstants(prob, 1, q));
1979566063dSJacob Faibussowitsch     PetscCall(PetscDSSetObjective(prob, 0, &f0_j_re));
1989566063dSJacob Faibussowitsch     PetscCall(DMPlexComputeIntegralFEM(plexe, XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, 0)], tt, NULL));
199a587d139SMark   } else tt[0] = 0;
200e0eea495SMark   J_re = -ctx->n_0 * ctx->v_0 * PetscRealPart(tt[0]);
2019566063dSJacob Faibussowitsch   PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, XsubArray)); // read only
2029566063dSJacob Faibussowitsch   PetscCall(PetscFree(XsubArray));
203a587d139SMark 
20490d163c9SMark Adams   if (rectx->use_spitzer_eta) {
20590d163c9SMark Adams     E = ctx->Ez = spit_eta * (rectx->j - J_re);
20690d163c9SMark Adams   } else {
20790d163c9SMark Adams     E        = ctx->Ez; /* keep real E */
20890d163c9SMark Adams     rectx->j = J;       /* cache */
20990d163c9SMark Adams   }
21090d163c9SMark Adams 
211e0eea495SMark   ratio = E / J / spit_eta;
2129371c9d4SSatish Balay   if (stepi > 10 && !rectx->use_spitzer_eta && ((old_ratio - ratio < 1.e-6))) {
21390d163c9SMark Adams     rectx->pulse_start     = time + 0.98 * dt;
214a587d139SMark     rectx->use_spitzer_eta = PETSC_TRUE;
215a587d139SMark   }
2169566063dSJacob Faibussowitsch   PetscCall(TSGetConvergedReason(ts, &reason));
2179566063dSJacob Faibussowitsch   PetscCall(TSGetConvergedReason(ts, &reason));
21890d163c9SMark Adams   if ((rectx->plotting) || stepi == 0 || reason || rectx->pulse_start == time + 0.98 * dt) {
2199371c9d4SSatish Balay     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,
2209371c9d4SSatish Balay                           (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));
22163a3b9bcSJacob Faibussowitsch     PetscCheck(rectx->pulse_start != (time + 0.98 * dt), PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "Spitzer complete ratio=%g", (double)ratio);
222930e68a5SMark Adams   }
223e0eea495SMark   old_ratio = ratio;
224e0eea495SMark   PetscFunctionReturn(0);
225e0eea495SMark }
226e0eea495SMark 
227e0eea495SMark static const double ppp = 2;
2289371c9d4SSatish Balay 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) {
229e0eea495SMark           LandauCtx      *ctx   = (LandauCtx *)constants;
230e0eea495SMark           REctx          *rectx = (REctx *)ctx->data;
231e0eea495SMark           PetscInt        ii    = rectx->idx, i;
232e0eea495SMark           const PetscReal kT_m  = ctx->k * ctx->thermal_temps[ii] / ctx->masses[ii]; /* kT/m */
233e0eea495SMark           const PetscReal n     = ctx->n[ii];
234e0eea495SMark           PetscReal       diff, f_maxwell, v2 = 0, theta = 2 * kT_m / (ctx->v_0 * ctx->v_0); /* theta = 2kT/mc^2 */
235e0eea495SMark           for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
236e0eea495SMark   f_maxwell = n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
237e0eea495SMark           diff      = 2. * PETSC_PI * x[0] * (PetscRealPart(u[ii]) - f_maxwell);
238e0eea495SMark           f0[0]     = PetscPowReal(diff, ppp);
239e0eea495SMark }
2409371c9d4SSatish Balay 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) {
241e0eea495SMark   LandauCtx      *ctx   = (LandauCtx *)constants;
242e0eea495SMark   REctx          *rectx = (REctx *)ctx->data;
243e0eea495SMark   PetscInt        ii    = rectx->idx, i;
244e0eea495SMark   const PetscReal kT_m  = ctx->k * ctx->thermal_temps[ii] / ctx->masses[ii]; /* kT/m */
245e0eea495SMark   const PetscReal n     = ctx->n[ii];
246e0eea495SMark   PetscReal       f_maxwell, v2 = 0, theta = 2 * kT_m / (ctx->v_0 * ctx->v_0); /* theta = 2kT/mc^2 */
247e0eea495SMark   for (i = 0; i < dim; ++i) v2 += x[i] * x[i];
248e0eea495SMark   f_maxwell = 2. * PETSC_PI * x[0] * n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta));
249e0eea495SMark   f0[0]     = PetscPowReal(f_maxwell, ppp);
250e0eea495SMark }
251e0eea495SMark 
252e0eea495SMark /*  */
2539371c9d4SSatish Balay static PetscErrorCode testStable(TS ts, Vec X, PetscInt stepi, PetscReal time, PetscBool islast, LandauCtx *ctx, REctx *rectx) {
254e0eea495SMark   PetscDS     prob;
255e0eea495SMark   Vec         X2;
256e0eea495SMark   PetscReal   ediff, idiff = 0, lpm0, lpm1 = 1;
257e0eea495SMark   PetscScalar tt[LANDAU_MAX_SPECIES];
2588a6f2e61SMark Adams   DM          dm, plex = ctx->plex[0];
259e0eea495SMark 
260e0eea495SMark   PetscFunctionBeginUser;
2619566063dSJacob Faibussowitsch   PetscCall(VecGetDM(X, &dm));
2629566063dSJacob Faibussowitsch   PetscCall(DMGetDS(plex, &prob));
2639566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(X, &X2));
2649566063dSJacob Faibussowitsch   PetscCall(VecCopy(X, X2));
265e0eea495SMark   if (!rectx->X_0) {
2669566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(X, &rectx->X_0));
2679566063dSJacob Faibussowitsch     PetscCall(VecCopy(X, rectx->X_0));
268e0eea495SMark   }
2699566063dSJacob Faibussowitsch   PetscCall(VecAXPY(X, -1.0, rectx->X_0));
2709566063dSJacob Faibussowitsch   PetscCall(PetscDSSetConstants(prob, sizeof(LandauCtx) / sizeof(PetscScalar), (PetscScalar *)ctx));
271e0eea495SMark   rectx->idx = 0;
2729566063dSJacob Faibussowitsch   PetscCall(PetscDSSetObjective(prob, 0, &f0_0_diff_lp));
2739566063dSJacob Faibussowitsch   PetscCall(DMPlexComputeIntegralFEM(plex, X2, tt, NULL));
274e0eea495SMark   ediff = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
2759566063dSJacob Faibussowitsch   PetscCall(PetscDSSetObjective(prob, 0, &f0_0_maxwellian_lp));
2769566063dSJacob Faibussowitsch   PetscCall(DMPlexComputeIntegralFEM(plex, X2, tt, NULL));
277e0eea495SMark   lpm0 = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
278e0eea495SMark   if (ctx->num_species > 1) {
279e0eea495SMark     rectx->idx = 1;
2809566063dSJacob Faibussowitsch     PetscCall(PetscDSSetObjective(prob, 0, &f0_0_diff_lp));
2819566063dSJacob Faibussowitsch     PetscCall(DMPlexComputeIntegralFEM(plex, X2, tt, NULL));
282e0eea495SMark     idiff = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
2839566063dSJacob Faibussowitsch     PetscCall(PetscDSSetObjective(prob, 0, &f0_0_maxwellian_lp));
2849566063dSJacob Faibussowitsch     PetscCall(DMPlexComputeIntegralFEM(plex, X2, tt, NULL));
285e0eea495SMark     lpm1 = PetscPowReal(PetscRealPart(tt[0]), 1. / ppp);
286e0eea495SMark   }
28763a3b9bcSJacob Faibussowitsch   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)));
288e0eea495SMark   /* view */
2899566063dSJacob Faibussowitsch   PetscCall(VecCopy(X2, X));
2909566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&X2));
291e0eea495SMark   if (islast) {
2929566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&rectx->X_0));
293e0eea495SMark     rectx->X_0 = NULL;
294e0eea495SMark   }
295e0eea495SMark   PetscFunctionReturn(0);
296e0eea495SMark }
297e0eea495SMark 
2989371c9d4SSatish Balay static PetscErrorCode EInduction(Vec X, Vec X_t, PetscInt step, PetscReal time, LandauCtx *ctx, PetscReal *a_E) {
299e0eea495SMark   REctx      *rectx = (REctx *)ctx->data;
300e0eea495SMark   PetscInt    ii;
301e0eea495SMark   DM          dm, plex;
3028a6f2e61SMark Adams   PetscScalar tt[LANDAU_MAX_SPECIES], qv0[LANDAU_MAX_SPECIES];
303e0eea495SMark   PetscReal   dJ_dt;
304e0eea495SMark   PetscDS     prob;
305e0eea495SMark 
306e0eea495SMark   PetscFunctionBeginUser;
3078a6f2e61SMark Adams   for (ii = 0; ii < ctx->num_species; ii++) qv0[ii] = ctx->charges[ii] * ctx->v_0;
3089566063dSJacob Faibussowitsch   PetscCall(VecGetDM(X, &dm));
3099566063dSJacob Faibussowitsch   PetscCall(DMGetDS(dm, &prob));
3109566063dSJacob Faibussowitsch   PetscCall(DMConvert(dm, DMPLEX, &plex));
311e0eea495SMark   /* get d current / dt */
3129566063dSJacob Faibussowitsch   PetscCall(PetscDSSetConstants(prob, ctx->num_species, qv0));
3139566063dSJacob Faibussowitsch   PetscCall(PetscDSSetObjective(prob, 0, &f0_jz_sum));
3143c633725SBarry Smith   PetscCheck(X_t, PETSC_COMM_SELF, PETSC_ERR_PLIB, "X_t");
3159566063dSJacob Faibussowitsch   PetscCall(DMPlexComputeIntegralFEM(plex, X_t, tt, NULL));
3168a6f2e61SMark Adams   dJ_dt = -ctx->n_0 * PetscRealPart(tt[0]) / ctx->t_0;
317e0eea495SMark   /* E induction */
318e0eea495SMark   *a_E  = -rectx->L * dJ_dt + rectx->Ez_initial;
3199566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&plex));
320e0eea495SMark   PetscFunctionReturn(0);
321e0eea495SMark }
322e0eea495SMark 
3239371c9d4SSatish Balay static PetscErrorCode EConstant(Vec X, Vec X_t, PetscInt step, PetscReal time, LandauCtx *ctx, PetscReal *a_E) {
324e0eea495SMark   PetscFunctionBeginUser;
325e0eea495SMark   *a_E = ctx->Ez;
326e0eea495SMark   PetscFunctionReturn(0);
327e0eea495SMark }
328e0eea495SMark 
3299371c9d4SSatish Balay static PetscErrorCode ENone(Vec X, Vec X_t, PetscInt step, PetscReal time, LandauCtx *ctx, PetscReal *a_E) {
330e0eea495SMark   PetscFunctionBeginUser;
331e0eea495SMark   *a_E = 0;
332e0eea495SMark   PetscFunctionReturn(0);
333e0eea495SMark }
334e0eea495SMark 
335e0eea495SMark /* ------------------------------------------------------------------- */
336e0eea495SMark /*
337e0eea495SMark    FormSource - Evaluates source terms F(t).
338e0eea495SMark 
339e0eea495SMark    Input Parameters:
340e0eea495SMark .  ts - the TS context
341e0eea495SMark .  time -
342e0eea495SMark .  X_dummmy - input vector
343e0eea495SMark .  dummy - optional user-defined context, as set by SNESSetFunction()
344e0eea495SMark 
345e0eea495SMark    Output Parameter:
346e0eea495SMark .  F - function vector
347e0eea495SMark  */
3489371c9d4SSatish Balay static PetscErrorCode FormSource(TS ts, PetscReal ftime, Vec X_dummmy, Vec F, void *dummy) {
349e0eea495SMark   PetscReal  new_imp_rate;
350e0eea495SMark   LandauCtx *ctx;
3518a6f2e61SMark Adams   DM         pack;
352e0eea495SMark   REctx     *rectx;
353e0eea495SMark 
354e0eea495SMark   PetscFunctionBeginUser;
3559566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &pack));
3569566063dSJacob Faibussowitsch   PetscCall(DMGetApplicationContext(pack, &ctx));
357e0eea495SMark   rectx = (REctx *)ctx->data;
358e0eea495SMark   /* check for impurities */
3599566063dSJacob Faibussowitsch   PetscCall(rectx->impuritySrcRate(ftime, &new_imp_rate, ctx));
360e0eea495SMark   if (new_imp_rate != 0) {
361e0eea495SMark     if (new_imp_rate != rectx->current_rate) {
362e0eea495SMark       PetscInt  ii;
363e0eea495SMark       PetscReal dne_dt, dni_dt, tilda_ns[LANDAU_MAX_SPECIES], temps[LANDAU_MAX_SPECIES];
3648fdabdddSMark Adams       Vec       globFarray[LANDAU_MAX_GRIDS * LANDAU_MAX_BATCH_SZ];
365e0eea495SMark       rectx->current_rate = new_imp_rate;
366e0eea495SMark       for (ii = 1; ii < LANDAU_MAX_SPECIES; ii++) tilda_ns[ii] = 0;
367e0eea495SMark       for (ii = 1; ii < LANDAU_MAX_SPECIES; ii++) temps[ii] = 1;
3688a6f2e61SMark Adams       dni_dt                   = new_imp_rate /* *ctx->t_0 */; /* fully ionized immediately, no normalize, stay in non-dim */
3698a6f2e61SMark Adams       dne_dt                   = new_imp_rate * rectx->Ne_ion /* *ctx->t_0 */;
3709371c9d4SSatish Balay       tilda_ns[0]              = dne_dt;
3719371c9d4SSatish Balay       tilda_ns[rectx->imp_idx] = dni_dt;
3729371c9d4SSatish Balay       temps[0]                 = rectx->T_cold;
3739371c9d4SSatish Balay       temps[rectx->imp_idx]    = rectx->T_cold;
37463a3b9bcSJacob Faibussowitsch       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));
3759566063dSJacob Faibussowitsch       PetscCall(DMCompositeGetAccessArray(pack, F, ctx->num_grids * ctx->batch_sz, NULL, globFarray));
3768a6f2e61SMark Adams       for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
377e0eea495SMark         /* add it */
378c3e4dd79SMark Adams         PetscCall(DMPlexLandauAddMaxwellians(ctx->plex[grid], globFarray[LAND_PACK_IDX(0, grid)], ftime, temps, tilda_ns, grid, 0, 1, ctx));
379e0eea495SMark       }
3808fdabdddSMark Adams       // Does DMCompositeRestoreAccessArray copy the data back? (no)
3819566063dSJacob Faibussowitsch       PetscCall(DMCompositeRestoreAccessArray(pack, F, ctx->num_grids * ctx->batch_sz, NULL, globFarray));
382e0eea495SMark     }
383e0eea495SMark   } else {
3849566063dSJacob Faibussowitsch     PetscCall(VecZeroEntries(F));
385e0eea495SMark     rectx->current_rate = 0;
386e0eea495SMark   }
387e0eea495SMark   PetscFunctionReturn(0);
388e0eea495SMark }
389419c2fb8Smarkadams4 
3909371c9d4SSatish Balay PetscErrorCode Monitor(TS ts, PetscInt stepi, PetscReal time, Vec X, void *actx) {
391e0eea495SMark   LandauCtx        *ctx   = (LandauCtx *)actx; /* user-defined application context */
392e0eea495SMark   REctx            *rectx = (REctx *)ctx->data;
393419c2fb8Smarkadams4   DM                pack  = NULL;
3948fdabdddSMark Adams   Vec               globXArray[LANDAU_MAX_GRIDS * LANDAU_MAX_BATCH_SZ];
395e0eea495SMark   TSConvergedReason reason;
396419c2fb8Smarkadams4 
397e0eea495SMark   PetscFunctionBeginUser;
398419c2fb8Smarkadams4   PetscCall(TSGetConvergedReason(ts, &reason));
399419c2fb8Smarkadams4   if (rectx->grid_view_idx != -1 || (reason && ctx->verbose > 3)) {
4009566063dSJacob Faibussowitsch     PetscCall(VecGetDM(X, &pack));
4019566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccessArray(pack, X, ctx->num_grids * ctx->batch_sz, NULL, globXArray));
402419c2fb8Smarkadams4   }
403e0eea495SMark   if (stepi > rectx->plotStep && rectx->plotting) {
404e0eea495SMark     rectx->plotting = PETSC_FALSE; /* was doing diagnostics, now done */
405e0eea495SMark     rectx->plotIdx++;
406e0eea495SMark   }
407e0eea495SMark   /* view */
408e0eea495SMark   if (time / rectx->plotDt >= (PetscReal)rectx->plotIdx || reason) {
409c3e4dd79SMark Adams     if ((reason || stepi == 0 || rectx->plotIdx % rectx->print_period == 0) && ctx->verbose > 1) {
410e0eea495SMark       /* print norms */
4119566063dSJacob Faibussowitsch       PetscCall(DMPlexLandauPrintNorms(X, stepi));
412a587d139SMark     }
413930e68a5SMark Adams     if (!rectx->plotting) { /* first step of possible backtracks */
414930e68a5SMark Adams       rectx->plotting = PETSC_TRUE;
415930e68a5SMark Adams       /* diagnostics + change E field with Sptizer (not just a monitor) */
4169566063dSJacob Faibussowitsch       PetscCall(rectx->test(ts, X, stepi, time, reason ? PETSC_TRUE : PETSC_FALSE, ctx, rectx));
417930e68a5SMark Adams     } else {
418930e68a5SMark Adams       PetscPrintf(PETSC_COMM_WORLD, "\t\t ERROR SKIP test spit ------\n");
419930e68a5SMark Adams       rectx->plotting = PETSC_TRUE;
420930e68a5SMark Adams     }
421419c2fb8Smarkadams4     if (rectx->grid_view_idx != -1) {
4229566063dSJacob Faibussowitsch       PetscCall(PetscObjectSetName((PetscObject)globXArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], rectx->grid_view_idx == 0 ? "ue" : "ui"));
423930e68a5SMark Adams       /* view, overwrite step when back tracked */
424e5d13be4Smarkadams4       PetscCall(DMSetOutputSequenceNumber(ctx->plex[rectx->grid_view_idx], rectx->plotIdx, time * ctx->t_0));
425419c2fb8Smarkadams4       PetscCall(VecViewFromOptions(globXArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], NULL, "-ex2_vec_view"));
426419c2fb8Smarkadams4     }
427e0eea495SMark     rectx->plotStep = stepi;
428930e68a5SMark Adams   } else {
42963a3b9bcSJacob Faibussowitsch     if (rectx->plotting) PetscPrintf(PETSC_COMM_WORLD, " ERROR rectx->plotting=%s step %" PetscInt_FMT "\n", PetscBools[rectx->plotting], stepi);
430930e68a5SMark Adams     /* diagnostics + change E field with Sptizer (not just a monitor) - can we lag this? */
4319566063dSJacob Faibussowitsch     PetscCall(rectx->test(ts, X, stepi, time, reason ? PETSC_TRUE : PETSC_FALSE, ctx, rectx));
432e0eea495SMark   }
433f37ccb5fSMark Adams   /* parallel check that only works of all batches are identical */
434f37ccb5fSMark Adams   if (reason && ctx->verbose > 3) {
435e0eea495SMark     PetscReal   val, rval;
436e0eea495SMark     PetscMPIInt rank;
4379566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
438f37ccb5fSMark Adams     for (PetscInt grid = 0; grid < ctx->num_grids; grid++) {
439f37ccb5fSMark Adams       PetscInt nerrors = 0;
440f37ccb5fSMark Adams       for (PetscInt i = 0; i < ctx->batch_sz; i++) {
4419566063dSJacob Faibussowitsch         PetscCall(VecNorm(globXArray[LAND_PACK_IDX(i, grid)], NORM_2, &val));
442f37ccb5fSMark Adams         if (i == 0) rval = val;
443f37ccb5fSMark Adams         else if ((val = PetscAbs(val - rval) / rval) > 1000 * PETSC_MACHINE_EPSILON) {
44463a3b9bcSJacob Faibussowitsch           PetscCall(PetscPrintf(PETSC_COMM_SELF, " [%d] Warning %" PetscInt_FMT ".%" PetscInt_FMT ") diff = %2.15e\n", rank, grid, i, (double)val));
445f37ccb5fSMark Adams           nerrors++;
446f37ccb5fSMark Adams         }
447f37ccb5fSMark Adams       }
448f37ccb5fSMark Adams       if (nerrors) {
44963a3b9bcSJacob Faibussowitsch         PetscCall(PetscPrintf(PETSC_COMM_SELF, " ***** [%d] ERROR max %" PetscInt_FMT " errors\n", rank, nerrors));
450e0eea495SMark       } else {
45163a3b9bcSJacob Faibussowitsch         PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[%d] %" PetscInt_FMT ") batch consistency check OK\n", rank, grid));
452f37ccb5fSMark Adams       }
453e0eea495SMark     }
454e0eea495SMark   }
455e0eea495SMark   rectx->idx = 0;
456*48a46eb9SPierre Jolivet   if (rectx->grid_view_idx != -1 || (reason && ctx->verbose > 3)) PetscCall(DMCompositeRestoreAccessArray(pack, X, ctx->num_grids * ctx->batch_sz, NULL, globXArray));
457e0eea495SMark   PetscFunctionReturn(0);
458e0eea495SMark }
459e0eea495SMark 
4609371c9d4SSatish Balay PetscErrorCode PreStep(TS ts) {
461e0eea495SMark   LandauCtx *ctx;
462e0eea495SMark   REctx     *rectx;
463e0eea495SMark   DM         dm;
464e0eea495SMark   PetscInt   stepi;
465e0eea495SMark   PetscReal  time;
466e0eea495SMark   Vec        X;
467e0eea495SMark 
468e0eea495SMark   PetscFunctionBeginUser;
469930e68a5SMark Adams   /* not used */
4709566063dSJacob Faibussowitsch   PetscCall(TSGetDM(ts, &dm));
4719566063dSJacob Faibussowitsch   PetscCall(TSGetTime(ts, &time));
4729566063dSJacob Faibussowitsch   PetscCall(TSGetSolution(ts, &X));
4739566063dSJacob Faibussowitsch   PetscCall(DMGetApplicationContext(dm, &ctx));
474e0eea495SMark   rectx = (REctx *)ctx->data;
4759566063dSJacob Faibussowitsch   PetscCall(TSGetStepNumber(ts, &stepi));
476e0eea495SMark   /* update E */
4779566063dSJacob Faibussowitsch   PetscCall(rectx->E(X, NULL, stepi, time, ctx, &ctx->Ez));
478e0eea495SMark   PetscFunctionReturn(0);
479e0eea495SMark }
480e0eea495SMark 
481e0eea495SMark /* 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) */
4829371c9d4SSatish Balay static PetscErrorCode stepSrc(PetscReal time, PetscReal *rho, LandauCtx *ctx) {
483e0eea495SMark   REctx *rectx = (REctx *)ctx->data;
484e0eea495SMark 
485e0eea495SMark   PetscFunctionBeginUser;
486e0eea495SMark   if (time >= rectx->pulse_start) *rho = rectx->pulse_rate;
487e0eea495SMark   else *rho = 0.;
488e0eea495SMark   PetscFunctionReturn(0);
489e0eea495SMark }
4909371c9d4SSatish Balay static PetscErrorCode zeroSrc(PetscReal time, PetscReal *rho, LandauCtx *ctx) {
491e0eea495SMark   PetscFunctionBeginUser;
492e0eea495SMark   *rho = 0.;
493e0eea495SMark   PetscFunctionReturn(0);
494e0eea495SMark }
4959371c9d4SSatish Balay static PetscErrorCode pulseSrc(PetscReal time, PetscReal *rho, LandauCtx *ctx) {
496e0eea495SMark   REctx *rectx = (REctx *)ctx->data;
497e0eea495SMark 
498e0eea495SMark   PetscFunctionBeginUser;
49908401ef6SPierre Jolivet   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'");
500e0eea495SMark   if (time < rectx->pulse_start || time > rectx->pulse_start + 3 * rectx->pulse_width) *rho = 0;
501a587d139SMark   else {
502e0eea495SMark     double x = PetscSinReal((time - rectx->pulse_start) / (3 * rectx->pulse_width) * 2 * PETSC_PI - PETSC_PI / 2) + 1; /* 0:2, integrates to 1.0 */
503e0eea495SMark     *rho     = rectx->pulse_rate * x / (3 * rectx->pulse_width);
504a587d139SMark     if (!rectx->use_spitzer_eta) rectx->use_spitzer_eta = PETSC_TRUE; /* use it next time */
505e0eea495SMark   }
506e0eea495SMark   PetscFunctionReturn(0);
507e0eea495SMark }
508e0eea495SMark 
509e0eea495SMark #undef __FUNCT__
510e0eea495SMark #define __FUNCT__ "ProcessREOptions"
5119371c9d4SSatish Balay static PetscErrorCode ProcessREOptions(REctx *rectx, const LandauCtx *ctx, DM dm, const char prefix[]) {
512e0eea495SMark   PetscFunctionList plist = NULL, testlist = NULL, elist = NULL;
513e0eea495SMark   char              pname[256], testname[256], ename[256];
514930e68a5SMark Adams   DM                dm_dummy;
515a587d139SMark   PetscBool         Connor_E = PETSC_FALSE;
516e0eea495SMark 
517e0eea495SMark   PetscFunctionBeginUser;
5189566063dSJacob Faibussowitsch   PetscCall(DMCreate(PETSC_COMM_WORLD, &dm_dummy));
519e0eea495SMark   rectx->Ne_ion          = 1;    /* number of electrons given up by impurity ion */
520e0eea495SMark   rectx->T_cold          = .005; /* kev */
521e0eea495SMark   rectx->ion_potential   = 15;   /* ev */
522e0eea495SMark   rectx->L               = 2;
523e0eea495SMark   rectx->X_0             = NULL;
524e0eea495SMark   rectx->imp_idx         = ctx->num_species - 1; /* default ionized impurity as last one */
525a587d139SMark   rectx->pulse_start     = PETSC_MAX_REAL;
526e0eea495SMark   rectx->pulse_width     = 1;
527e0eea495SMark   rectx->plotStep        = PETSC_MAX_INT;
528e0eea495SMark   rectx->pulse_rate      = 1.e-1;
529e0eea495SMark   rectx->current_rate    = 0;
530e0eea495SMark   rectx->plotIdx         = 0;
531e0eea495SMark   rectx->j               = 0;
532e0eea495SMark   rectx->plotDt          = 1.0;
533930e68a5SMark Adams   rectx->plotting        = PETSC_FALSE;
534e0eea495SMark   rectx->use_spitzer_eta = PETSC_FALSE;
535e0eea495SMark   rectx->idx             = 0;
5368a6f2e61SMark Adams   rectx->print_period    = 10;
537419c2fb8Smarkadams4   rectx->grid_view_idx   = -1; // do not get if not needed
538e0eea495SMark   /* Register the available impurity sources */
5399566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&plist, "step", &stepSrc));
5409566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&plist, "none", &zeroSrc));
5419566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&plist, "pulse", &pulseSrc));
5429566063dSJacob Faibussowitsch   PetscCall(PetscStrcpy(pname, "none"));
5439566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&testlist, "none", &testNone));
5449566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&testlist, "spitzer", &testSpitzer));
5459566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&testlist, "stable", &testStable));
5469566063dSJacob Faibussowitsch   PetscCall(PetscStrcpy(testname, "none"));
5479566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&elist, "none", &ENone));
5489566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&elist, "induction", &EInduction));
5499566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListAdd(&elist, "constant", &EConstant));
5509566063dSJacob Faibussowitsch   PetscCall(PetscStrcpy(ename, "constant"));
551e0eea495SMark 
552d0609cedSBarry Smith   PetscOptionsBegin(PETSC_COMM_SELF, prefix, "Options for Runaway/seed electron model", "none");
5539566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-ex2_plot_dt", "Plotting interval", "ex2.c", rectx->plotDt, &rectx->plotDt, NULL));
554e0eea495SMark   if (rectx->plotDt < 0) rectx->plotDt = 1e30;
555e0eea495SMark   if (rectx->plotDt == 0) rectx->plotDt = 1e-30;
5569566063dSJacob Faibussowitsch   PetscCall(PetscOptionsInt("-ex2_print_period", "Plotting interval", "ex2.c", rectx->print_period, &rectx->print_period, NULL));
5579566063dSJacob Faibussowitsch   PetscCall(PetscOptionsInt("-ex2_grid_view_idx", "grid_view_idx", "ex2.c", rectx->grid_view_idx, &rectx->grid_view_idx, NULL));
558419c2fb8Smarkadams4   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);
5599566063dSJacob Faibussowitsch   PetscCall(PetscOptionsFList("-ex2_impurity_source_type", "Name of impurity source to run", "", plist, pname, pname, sizeof(pname), NULL));
5609566063dSJacob Faibussowitsch   PetscCall(PetscOptionsFList("-ex2_test_type", "Name of test to run", "", testlist, testname, testname, sizeof(testname), NULL));
5619566063dSJacob Faibussowitsch   PetscCall(PetscOptionsInt("-ex2_impurity_index", "index of sink for impurities", "none", rectx->imp_idx, &rectx->imp_idx, NULL));
562cad9d221SBarry Smith   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);
5639566063dSJacob Faibussowitsch   PetscCall(PetscOptionsFList("-ex2_e_field_type", "Electric field type", "", elist, ename, ename, sizeof(ename), NULL));
564e0eea495SMark   rectx->Ne_ion = -ctx->charges[rectx->imp_idx] / ctx->charges[0];
5659566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-ex2_t_cold", "Temperature of cold electron and ions after ionization in keV", "none", rectx->T_cold, &rectx->T_cold, NULL));
5669566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-ex2_pulse_start_time", "Time at which pulse happens for 'pulse' source", "none", rectx->pulse_start, &rectx->pulse_start, NULL));
5679566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-ex2_pulse_width_time", "Width of pulse 'pulse' source", "none", rectx->pulse_width, &rectx->pulse_width, NULL));
5689566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-ex2_pulse_rate", "Number density of pulse for 'pulse' source", "none", rectx->pulse_rate, &rectx->pulse_rate, NULL));
569e0eea495SMark   rectx->T_cold *= 1.16e7; /* convert to Kelvin */
5709566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-ex2_ion_potential", "Potential to ionize impurity (should be array) in ev", "none", rectx->ion_potential, &rectx->ion_potential, NULL));
5719566063dSJacob Faibussowitsch   PetscCall(PetscOptionsReal("-ex2_inductance", "Inductance E feild", "none", rectx->L, &rectx->L, NULL));
5729566063dSJacob Faibussowitsch   PetscCall(PetscOptionsBool("-ex2_connor_e_field_units", "Scale Ex but Connor-Hastie E_c", "none", Connor_E, &Connor_E, NULL));
57363a3b9bcSJacob Faibussowitsch   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));
574d0609cedSBarry Smith   PetscOptionsEnd();
575e0eea495SMark   /* get impurity source rate function */
5769566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListFind(plist, pname, &rectx->impuritySrcRate));
5773c633725SBarry Smith   PetscCheck(rectx->impuritySrcRate, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "No impurity source function found '%s'", pname);
5789566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListFind(testlist, testname, &rectx->test));
5793c633725SBarry Smith   PetscCheck(rectx->test, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "No test found '%s'", testname);
5809566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListFind(elist, ename, &rectx->E));
5813c633725SBarry Smith   PetscCheck(rectx->E, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "No E field function found '%s'", ename);
5829566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListDestroy(&plist));
5839566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListDestroy(&testlist));
5849566063dSJacob Faibussowitsch   PetscCall(PetscFunctionListDestroy(&elist));
585930e68a5SMark Adams 
586a587d139SMark   /* convert E from Connor-Hastie E_c units to real if doing Spitzer E */
587a587d139SMark   if (Connor_E) {
588e0eea495SMark     PetscReal E = ctx->Ez, Tev = ctx->thermal_temps[0] * 8.621738e-5, n = ctx->n_0 * ctx->n[0];
589e0eea495SMark     CalculateE(Tev, n, ctx->lnLam, ctx->epsilon0, &E);
590e0eea495SMark     ((LandauCtx *)ctx)->Ez *= E;
591e0eea495SMark   }
5929566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&dm_dummy));
593e0eea495SMark   PetscFunctionReturn(0);
594e0eea495SMark }
595e0eea495SMark 
5969371c9d4SSatish Balay int main(int argc, char **argv) {
5978a6f2e61SMark Adams   DM         pack;
598419c2fb8Smarkadams4   Vec        X;
599f37ccb5fSMark Adams   PetscInt   dim = 2, nDMs;
600e0eea495SMark   TS         ts;
601e0eea495SMark   Mat        J;
602e0eea495SMark   PetscDS    prob;
603e0eea495SMark   LandauCtx *ctx;
604e0eea495SMark   REctx     *rectx;
605a587d139SMark #if defined PETSC_USE_LOG
606a587d139SMark   PetscLogStage stage;
607a587d139SMark #endif
608930e68a5SMark Adams   PetscMPIInt rank;
6098fdabdddSMark Adams #if defined(PETSC_HAVE_THREADSAFETY)
6108fdabdddSMark Adams   double starttime, endtime;
6118fdabdddSMark Adams #endif
612327415f7SBarry Smith   PetscFunctionBeginUser;
6139566063dSJacob Faibussowitsch   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
6149566063dSJacob Faibussowitsch   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
615930e68a5SMark Adams   if (rank) { /* turn off output stuff for duplicate runs */
616419c2fb8Smarkadams4     PetscCall(PetscOptionsClearValue(NULL, "-ex2_dm_view"));
617419c2fb8Smarkadams4     PetscCall(PetscOptionsClearValue(NULL, "-ex2_vec_view"));
618419c2fb8Smarkadams4     PetscCall(PetscOptionsClearValue(NULL, "-ex2_vec_view_init"));
619419c2fb8Smarkadams4     PetscCall(PetscOptionsClearValue(NULL, "-ex2_dm_view_init"));
6209566063dSJacob Faibussowitsch     PetscCall(PetscOptionsClearValue(NULL, "-info")); /* this does not work */
621930e68a5SMark Adams   }
6229566063dSJacob Faibussowitsch   PetscCall(PetscOptionsGetInt(NULL, NULL, "-dim", &dim, NULL));
623e0eea495SMark   /* Create a mesh */
6249566063dSJacob Faibussowitsch   PetscCall(DMPlexLandauCreateVelocitySpace(PETSC_COMM_WORLD, dim, "", &X, &J, &pack));
6259566063dSJacob Faibussowitsch   PetscCall(DMCompositeGetNumberDM(pack, &nDMs));
6269566063dSJacob Faibussowitsch   PetscCall(PetscObjectSetName((PetscObject)J, "Jacobian"));
6279566063dSJacob Faibussowitsch   PetscCall(PetscObjectSetName((PetscObject)X, "f"));
6289566063dSJacob Faibussowitsch   PetscCall(DMGetApplicationContext(pack, &ctx));
6299566063dSJacob Faibussowitsch   PetscCall(DMSetUp(pack));
630e0eea495SMark   /* context */
6319566063dSJacob Faibussowitsch   PetscCall(PetscNew(&rectx));
632e0eea495SMark   ctx->data = rectx;
6339566063dSJacob Faibussowitsch   PetscCall(ProcessREOptions(rectx, ctx, pack, ""));
6349566063dSJacob Faibussowitsch   PetscCall(DMGetDS(pack, &prob));
635419c2fb8Smarkadams4   if (rectx->grid_view_idx != -1) {
636419c2fb8Smarkadams4     Vec *XsubArray = NULL;
637419c2fb8Smarkadams4     PetscCall(PetscMalloc(sizeof(*XsubArray) * nDMs, &XsubArray));
6389566063dSJacob Faibussowitsch     PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, XsubArray)); // read only
6399566063dSJacob Faibussowitsch     PetscCall(PetscObjectSetName((PetscObject)XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], rectx->grid_view_idx == 0 ? "ue" : "ui"));
640e5d13be4Smarkadams4     PetscCall(DMSetOutputSequenceNumber(ctx->plex[rectx->grid_view_idx], 0, 0.0));
641419c2fb8Smarkadams4     PetscCall(DMViewFromOptions(ctx->plex[rectx->grid_view_idx], NULL, "-ex2_dm_view"));
642419c2fb8Smarkadams4     PetscCall(DMViewFromOptions(ctx->plex[rectx->grid_view_idx], NULL, "-ex2_dm_view_init"));
643e5d13be4Smarkadams4     PetscCall(VecViewFromOptions(XsubArray[LAND_PACK_IDX(ctx->batch_view_idx, rectx->grid_view_idx)], NULL, "-ex2_vec_view"));      // initial condition (monitor plots after step)
644419c2fb8Smarkadams4     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)
6459566063dSJacob Faibussowitsch     PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, XsubArray));                                                       // read only
6469566063dSJacob Faibussowitsch     PetscCall(PetscFree(XsubArray));
647419c2fb8Smarkadams4   }
648e0eea495SMark   /* Create timestepping solver context */
6499566063dSJacob Faibussowitsch   PetscCall(TSCreate(PETSC_COMM_SELF, &ts));
6509566063dSJacob Faibussowitsch   PetscCall(TSSetDM(ts, pack));
6519566063dSJacob Faibussowitsch   PetscCall(TSSetIFunction(ts, NULL, DMPlexLandauIFunction, NULL));
6529566063dSJacob Faibussowitsch   PetscCall(TSSetIJacobian(ts, J, J, DMPlexLandauIJacobian, NULL));
6539566063dSJacob Faibussowitsch   PetscCall(TSSetRHSFunction(ts, NULL, FormSource, NULL));
6549566063dSJacob Faibussowitsch   PetscCall(TSSetFromOptions(ts));
6559566063dSJacob Faibussowitsch   PetscCall(TSSetSolution(ts, X));
6569566063dSJacob Faibussowitsch   PetscCall(TSSetApplicationContext(ts, ctx));
6579566063dSJacob Faibussowitsch   PetscCall(TSMonitorSet(ts, Monitor, ctx, NULL));
6589566063dSJacob Faibussowitsch   PetscCall(TSSetPreStep(ts, PreStep));
659e0eea495SMark   rectx->Ez_initial = ctx->Ez; /* cache for induction caclulation - applied E field */
6608594ddcfSMark Adams   if (1) {                     /* warm up an test just DMPlexLandauIJacobian */
661e0eea495SMark     Vec       vec;
662a587d139SMark     PetscInt  nsteps;
663a587d139SMark     PetscReal dt;
6649566063dSJacob Faibussowitsch     PetscCall(PetscLogStageRegister("Warmup", &stage));
6659566063dSJacob Faibussowitsch     PetscCall(PetscLogStagePush(stage));
6669566063dSJacob Faibussowitsch     PetscCall(VecDuplicate(X, &vec));
6679566063dSJacob Faibussowitsch     PetscCall(VecCopy(X, vec));
6689566063dSJacob Faibussowitsch     PetscCall(TSGetMaxSteps(ts, &nsteps));
6699566063dSJacob Faibussowitsch     PetscCall(TSGetTimeStep(ts, &dt));
6709566063dSJacob Faibussowitsch     PetscCall(TSSetMaxSteps(ts, 1));
6719566063dSJacob Faibussowitsch     PetscCall(TSSolve(ts, X));
6729566063dSJacob Faibussowitsch     PetscCall(TSSetMaxSteps(ts, nsteps));
6739566063dSJacob Faibussowitsch     PetscCall(TSSetStepNumber(ts, 0));
6749566063dSJacob Faibussowitsch     PetscCall(TSSetTime(ts, 0));
6759566063dSJacob Faibussowitsch     PetscCall(TSSetTimeStep(ts, dt));
676a587d139SMark     rectx->plotIdx  = 0;
677930e68a5SMark Adams     rectx->plotting = PETSC_FALSE;
6789566063dSJacob Faibussowitsch     PetscCall(PetscLogStagePop());
6799566063dSJacob Faibussowitsch     PetscCall(VecCopy(vec, X));
6809566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&vec));
681c3e4dd79SMark Adams     PetscCall(PetscObjectStateIncrease((PetscObject)ctx->J));
682e0eea495SMark   }
683e0eea495SMark   /* go */
6849566063dSJacob Faibussowitsch   PetscCall(PetscLogStageRegister("Solve", &stage));
685f37ccb5fSMark Adams   ctx->stage = 0; // lets not use this stage
6868fdabdddSMark Adams #if defined(PETSC_HAVE_THREADSAFETY)
687e607c864SMark Adams   ctx->stage = 1; // not set with thread safety
6888fdabdddSMark Adams #endif
689c3e4dd79SMark Adams   //PetscCall(TSSetSolution(ts,X));
6909566063dSJacob Faibussowitsch   PetscCall(PetscLogStagePush(stage));
6918fdabdddSMark Adams #if defined(PETSC_HAVE_THREADSAFETY)
6928fdabdddSMark Adams   starttime = MPI_Wtime();
6938fdabdddSMark Adams #endif
694419c2fb8Smarkadams4 #if defined(PETSC_HAVE_CUDA_NVTX)
695419c2fb8Smarkadams4   nvtxRangePushA("ex2-TSSolve-warm");
696419c2fb8Smarkadams4 #endif
6979566063dSJacob Faibussowitsch   PetscCall(TSSolve(ts, X));
698419c2fb8Smarkadams4 #if defined(PETSC_HAVE_CUDA_NVTX)
699419c2fb8Smarkadams4   nvtxRangePop();
700419c2fb8Smarkadams4 #endif
7019566063dSJacob Faibussowitsch   PetscCall(PetscLogStagePop());
7028fdabdddSMark Adams #if defined(PETSC_HAVE_THREADSAFETY)
7038fdabdddSMark Adams   endtime = MPI_Wtime();
7048fdabdddSMark Adams   ctx->times[LANDAU_EX2_TSSOLVE] += (endtime - starttime);
7058fdabdddSMark Adams #endif
706e0eea495SMark   /* clean up */
7079566063dSJacob Faibussowitsch   PetscCall(DMPlexLandauDestroyVelocitySpace(&pack));
7089566063dSJacob Faibussowitsch   PetscCall(TSDestroy(&ts));
7099566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&X));
7109566063dSJacob Faibussowitsch   PetscCall(PetscFree(rectx));
7119566063dSJacob Faibussowitsch   PetscCall(PetscFinalize());
712b122ec5aSJacob Faibussowitsch   return 0;
713e0eea495SMark }
714e0eea495SMark 
715e0eea495SMark /*TEST
716a587d139SMark 
7175dac466eSMark Adams   testset:
718984ed092SMark Adams     requires: p4est !complex double defined(PETSC_USE_DMLANDAU_2D)
7195dac466eSMark Adams     output_file: output/ex2_0.out
720c3e4dd79SMark Adams     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
7215dac466eSMark Adams     test:
7225dac466eSMark Adams       suffix: cpu
723c3e4dd79SMark Adams       args: -dm_landau_device_type cpu -ksp_type bicg -pc_type jacobi
724a587d139SMark     test:
725a587d139SMark       suffix: kokkos
7265dac466eSMark Adams       requires: kokkos_kernels
727c3e4dd79SMark Adams       args: -dm_landau_device_type kokkos -dm_mat_type aijkokkos -dm_vec_type kokkos -ksp_type bicg -pc_type jacobi
72890d163c9SMark Adams     test:
72990d163c9SMark Adams       suffix: cuda
7305dac466eSMark Adams       requires: cuda
731c3e4dd79SMark Adams       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
732cb25d741SMark Adams     test:
733cb25d741SMark Adams       suffix: kokkos_batch
734cb25d741SMark Adams       requires: kokkos_kernels
735c3e4dd79SMark Adams       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
736bfc784b7SMark Adams     test:
737bfc784b7SMark Adams       suffix: kokkos_batch_coo
738bfc784b7SMark Adams       requires: kokkos_kernels
739c3e4dd79SMark Adams       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
74090d163c9SMark Adams 
741e0eea495SMark TEST*/
742