Coupling with an AI model for sheath BC in GK simulations

[Antoinehoff]

This PR follows DR #957 and introduces $\mu$-dependent sheath boundary conditions in the gyrokinetic solver using a neural network surrogate of the GYRAZE code.

Motivation

The standard conducting sheath BC in Gkeyll uses a constant cutting velocity $v_{\parallel\text{cut}}$ determined by the potential drop $\Delta\phi = \phi_\text{mpe} - \phi_w$ between the magnetic presheath entrance (MPE) and the wall:

$$\frac{1}{2} m_s v_{\parallel\text{cut}}^2 = q_s \Delta\phi$$

This ignores the effect of the particle magnetic moment $\mu$ on the absorption/reflection condition.

New files

• gyrokinetic/ker/bc_sheath_gyrokinetic/bc_sheath_gyrokinetic_gyraze_surrogate.c
  Auto-generated GYRAZE surrogate model (NN weights, SVM classifier, forward pass, interpolation, projection), from gkeyll_sheath_ai, the commit is indicated in the file.

• gyrokinetic/ker/bc_sheath_gyrokinetic/gkyl_bc_sheath_gyrokinetic_gyraze_surrogate.h
  Public API for the surrogate.

• gyrokinetic/creg/rt_gk_tcv_sheath_surrogate_2x2v_p1.c
  Regression test, TCV clopen IWL GK simulation with surrogate sheath BC.

Modified files

• gyrokinetic/zero/gkyl_bc_sheath_gyrokinetic.h
  New use_surrogate parameter in constructor; new gkyl_bc_sheath_gyrokinetic_update_vcut_fact_surrogate and gkyl_bc_sheath_gyrokinetic_evaluate_vcut_fact_surrogate public methods; getters/setters for vcut_fact, its basis and range.

• gyrokinetic/zero/gkyl_bc_sheath_gyrokinetic_priv.h
gkyl_bc_sheath_gyrokinetic struct now stores: vcut_fact DG array over perpendicular config space $\times,\mu$, its basis and range, use_surrogate flag, function pointer update_vcut_fact for enabled/disabled dispatch, and surrogate_eval interface.

• gyrokinetic/zero/bc_sheath_gyrokinetic.c
  Initializes vcut_fact to 1 (reproducing standard BC); sets up surrogate kernel selection and function pointer dispatch; CPU loop for surrogate evaluation over the vcut_fact range.

• gyrokinetic/zero/bc_sheath_gyrokinetic_cu.cu
  CUDA kernel bc_gksheath_update_vcut_fact_surrogate_cu_ker parallelizes surrogate evaluation over the vcut_fact range on GPU.

• gyrokinetic/ker/bc_sheath_gyrokinetic/bc_sheath_gyrokinetic_ser_p1.c
  Reflection kernels now accept a vcut_fact DG expansion; new bc_sheath_gyrokinetic_surrogate_{lower,upper}_{1x2v,2x2v,3x2v}_ser_p1 kernels that extract physical quantities from DG coefficients, evaluate the surrogate, and write vcut_fact (generated by gkylcas#97).

• gyrokinetic/apps/gkyl_gyrokinetic.h
  New use_sheath_surrogate field in gkyl_gyrokinetic_bc.

• gyrokinetic/apps/gkyl_gyrokinetic_priv.h
  Species struct gains alloc_srg_aux_var, maxwell_mom, maxmom, dens_sheath, temp_sheath for surrogate moment computation.

• gyrokinetic/apps/gk_species.c
  Allocates Maxwellian moment calculator when surrogate is enabled; computes density and temperature before each BC application; calls gkyl_bc_sheath_gyrokinetic_update_vcut_fact_surrogate at both lower and upper sheath/IWL boundaries.

• gyrokinetic/zero/calc_metric.c, calc_metric_mirror.c
  Computes $\sin\alpha = 1/\sqrt{g_{33},g^{33}}$ and stores it in bimpactangle.

• gyrokinetic/unit/ctest_bc_sheath_gyrokinetic.c Expanded with surrogate tests: verifies vcut_fact evaluation, the direct surrogate interface, and end-to-end BC application with surrogate for 1x2v, 2x2v, and 3x2v on both CPU and GPU.

Design choices

Correcting factor approach

Following DR #957, a $\mu$-dependent multiplying factor $\alpha(\mu)$ modifies the cutting velocity:

$$v_{\parallel\text{cut}}^2(\mu) = \alpha(\mu) ,\frac{2,q_s}{m_s},\Delta\phi$$

The vcut_fact array stores $\alpha(\mu)$ as a DG expansion over perpendicular configuration space $\times,\mu$ (serendipity basis, same polynomial order as the phase-space basis). It is initialized to 1 everywhere, which recovers the standard conducting sheath BC. When the surrogate is enabled, vcut_fact is updated from the NN prediction before each BC application.

Surrogate model interface

The surrogate interface consists of two components auto-exported from Python to self-contained C (see gyraze_surrogate_c_interface repository).
All surrogate functions are marked GKYL_CU_DH and use only stack memory, making them compatible with both CPU and GPU execution without any external dependencies beyond libm.

The surrogate interface is designed with the same philosophy as the kernels, where an external code is used to generate the source files. One can easily update the surrogate by using the gyraze_surrogate_c_interface repository.

Two-step update

The vcut_fact update and the BC application are deliberately separated into two steps:

1. gkyl_bc_sheath_gyrokinetic_update_vcut_fact_surrogate(...) — evaluates the surrogate and writes vcut_fact.
2. gkyl_bc_sheath_gyrokinetic_advance(...) — applies the reflection BC using the stored vcut_fact.

This split allows future optimization where the surrogate is evaluated less frequently than the BC if evaluation cost becomes significant.

Safety assertions

• An assert enforces that the surrogate is only used for electrons (q2Dm == -2e/m_e), since ion reflection is not expected.
• The surrogate requires vdim > 1 (i.e. a $\mu$ dimension must exist).

User interface

Enabling the surrogate requires a single flag in the BC specification:

.bc_z = {
  { .dir = 1, 
    .edge = GKYL_LOWER_EDGE, 
    .type = GKYL_BC_GK_SPECIES_IWL,
    .use_sheath_surrogate = true 
  },
  { .dir = 1, 
    .edge = GKYL_UPPER_EDGE, 
    .type = GKYL_BC_GK_SPECIES_IWL, 
    .use_sheath_surrogate = true },
},

Demonstration

Unit test

The expanded unit test gyrokinetic/unit/ctest_bc_sheath_gyrokinetic.c validates:

• Standard BC application (without surrogate) for 1x2v, 2x2v, 3x2v at both edges, on CPU and GPU.
• Surrogate vcut_fact evaluation and end-to-end BC application with the surrogate.
• Direct surrogate interface (gkyl_bc_sheath_gyrokinetic_evaluate_vcut_fact_surrogate).

Below are outputs with write_field = true, showing the reflected distribution and the $v_\text{cut}(\mu)$ curve from the surrogate:

The $v_\text{cut}(\mu)$ curve is now provided by the GYRAZE surrogate rather than the ad-hoc function in #960.

TCV 2x2v regression test

The new regression test rt_gk_tcv_sheath_surrogate_2x2v_p1.c runs a TCV clopen IWL GK simulation. At higher resolution (24 $\times$ 16 $\times$ 12 $\times$ 8) we obtain the following results with and without the surrogate:

Without surrogate	With surrogate

We can see that the electron distribution functions close to the magnetic presheath entrance ($z=\pm \pi$) is affected by the surrogate as expected. More electrons are absorbed with the surrogate which creates a large region of $f_e=0$ in the $v_\parallel > 0$ region for $z=-\pi$ and $v_\parallel < 0$ region for $z=\pi$. This also show that the surrogate is active in the simulation and is not always reaching the non converged state which would bypass it. The plots are zoomed in the region of interest.

Without surrogate	With surrogate

The SOL potential is significantly higher with the surrogate. This is expected: FLR effects allow electrons to be absorbed at a gyroradius distance from the wall, lowering the effective potential barrier. The increased negative charge depletion raises the potential.

TCV 3x2v production like turbulence simulation

We consider now my favorite 3x2v low cost production case (see input file).

Without surrogate	With surrogate

We also check the distribution function at the magnetic presheath entrance for the 3x2v case. The effect of the surrogate is less visible on this snapshot. It is possible that the surrogate is often bypassed though an effect is observed in the potential and macroscopic quantities.

Without surrogate	With surrogate

Performance comparison

For production like run as the 3x2v simulation presented above, the computational cost is limited. Averaged over a 6h restart, we have a cost of forward Euler evaluation of:

• 54.7547 ms for CSBC (see csbc_timing.txt)
• 56.4295 ms for AIBC (see aibc_timing.txt)

Known limitations

• Small impact angle assumption: The GYRAZE code (and hence the surrogate) assumes a small magnetic field impact angle with the wall. The surrogate should not be used for configurations like LAPD where $\alpha \sim \pi/2$. A runtime check on the impact angle could be added.
• Training domain: The surrogate is trained on $\alpha \in [2°, 10°]$, $\gamma \in [0.5, 4]$, $\hat\phi \in [1, 10]$. Outside this domain, the SVM classifier will avoid the evaluation of the NN and return vcut_fact = 1.
• Electron-only: Currently restricted to electrons by assertion. There is no need of a surrogate for ions since they are mostly absorbed in the applications of interest, but this could be relaxed in the future if needed.