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.

It is now using the KANN Gkeyll interface developed in #1056. In particular, this branch demonstrates the use of KANN model inference on multi GPU in production simulations.

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.

The main changes in this PR reside in a new gkyl_array vcutsq which stores the $v_{\parallel\text{cut}}(\mu)$ curve in the gkyl_bc_sheath_gyrokinetic structure. The evaluation of the sheath BC is now done in two steps:

1. update of vcutsq, if we don't use the NN model, it will just contain the conducting sheath constant value.
2. set the ghost values with the reflection algorithm according to vcutsq.

The inference of the NN model in step 1. above is done in three stages:

1. Fill the KANN input vector with all parameter triplet $\alpha,\gamma,\hat\phi$ evaluated on each perpendicular surface nodes.
2. Call gkyl_kann_net_apply storing the output in kann_out. The KANN Gkeyll interface (GPU Hackathon 2026 branch: porting the KANN implementation to GPUs #1056) makes this call on CPU or GPU.
3. Convert back the nodal values produced by the NN to a DG representation in vcutsq. Note, this step requires to interpolate between the NN model hardcoded 20 points $\mu$ grid and the simulation $\mu$ grid. The grid points are expected to be output by the NN with an extension of the output layer so we do not have any hardcoded features now.

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. This is now replaced by the KANN Gkeyll interface.

• gyrokinetic/ker/bc_sheath_gyrokinetic/gkyl_bc_sheath_gyrokinetic_gyraze_surrogate.h
  Public API for the surrogate. This is now replaced by the KANN Gkeyll interface.

• gyrokinetic/creg/rt_gk_tcv_sheath_surrogate_2x2v_p1.c
  Regression test, TCV clopen IWL GK simulation with surrogate sheath BC. The original 2x2v tcv reg test is using the surrogate model now.

• gyrokinetic/data/nn_model/nn_model_sheath_bc_conv_MPE.kann
  This is the NN model in a KANN readable format, generate using a python routine to convert from pytorch. Nomenclature: "conv" indicates that the model was trained on converged region of GYRAZE, "MPE" means it is expecting values at the magnetic presheath entrance. One could think of simplifying this name...

Modified files

• gyrokinetic/zero/gkyl_bc_sheath_gyrokinetic.h
  New constructor parameters use_surrogate, surrogate_model_path, and phase_grid/phase_global (the latter two for the vcutsq diagnostic output). New public methods gkyl_bc_sheath_gyrokinetic_update_vcutsq (computes vcutsq), gkyl_bc_sheath_gyrokinetic_set_vcutsq, and gkyl_bc_sheath_gyrokinetic_write_vcutsq (writes vcutsq as a diagnostic).

• gyrokinetic/zero/gkyl_bc_sheath_gyrokinetic_priv.h
gkyl_bc_sheath_gyrokinetic struct now stores: vcutsq DG array over perpendicular config space $\times,\mu$, its basis and range, use_surrogate flag, a update_vcutsq function pointer for switching between conducting sheath and NN inference, and the gkyl_kn_vec input/output buffers and gkyl_kann_net of the KANN interface.

• gyrokinetic/zero/bc_sheath_gyrokinetic.c
  See git diff.

• gyrokinetic/zero/bc_sheath_gyrokinetic_cu.cu
  See git diff.

• gyrokinetic/ker/bc_sheath_gyrokinetic/bc_sheath_gyrokinetic_ser_p1.c
  Reflection kernels now take the vcutsq DG expansion and do not compute it on the fly anymore. New bc_sheath_gyrokinetic_vcutsq_const_{lower,upper}_{1x2v,2x2v,3x2v}_ser_p1 kernels to compute the constant conducting sheath model. New bc_sheath_gyrokinetic_vcutsq_surr_{lower,upper}_{1x2v,2x2v,3x2v}_ser_p1 kernels to compute the surrogate vcut.

• gyrokinetic/apps/gkyl_gyrokinetic.h
  New use_sheath_surrogate boolean to activate surrogate. New surrogate_model_path to store the path to the .kann file.

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

• gyrokinetic/apps/gk_species.c
  Allocates the Maxwellian moment calculator when the surrogate is enabled, computes density and temperature before each BC application, and calls gkyl_bc_sheath_gyrokinetic_update_vcutsq at both lower and upper sheath/IWL boundaries. Also writes the vcutsq diagnostic alongside the phase-space data.

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

Design choices

Correcting factor approach

Unlike DR #957, we gave up on the $\mu$-dependent multiplying factor $\alpha(\mu)$ which would modify the cutting velocity:

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

because of computing it requires to divide by $\Delta\phi$ that can be 0.

Instead, the vcutsq array stores the full cutting velocity as a DG expansion over perpendicular configuration space $\times,\mu$.

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,
    .surrogate_model_path = "gyrokinetic/data/nn_model/nn_model_sheath_bc_conv_MPE.kann",
  },
  },
  { .dir = 1, 
    .edge = GKYL_UPPER_EDGE, 
    .type = GKYL_BC_GK_SPECIES_IWL, 
    .use_sheath_surrogate = true,
    .surrogate_model_path = "gyrokinetic/data/nn_model/nn_model_sheath_bc_conv_MPE.kann",
  },
},

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 vcutsq evaluation and end-to-end BC application with the 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 c onfigurations 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, we still evaluate the surrogate but we don't know how physical it is.
• 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.

[manauref]

Very nice general description. It covers all the basis and offers useful plots. I also very much like that this is now using the kann code.

Here's some feedback on what I've had time to review. I'll try to look at the actual kann part later today/tomorrow. In addition to the minor things below, here are some higher level questions:

1. One downside I see here is that another person (maybe even yourself) could not extend the surrogate because we don't have enough material to replicate the construction of the network. Would it be possible to create a directory (say in gyrokinetic/data/) with the GYRAZE input files and network training scripts? I see this as necessary, much in the same way that we have a folder for creating the parameters of the radiation operator (which depend on other external analysis, but which can be replicated especially with what's in PR #1060).
2. I'd like to read your AI sheath paper draft before we merge this to better understand what this is contributing and how it's changing the physics (beyond what you've put in the description).
3. The description mentions that this model doesn't work for ion reflection. As you know Gkeyll can produce transients with negative MPE potential. What does the surrogate do to electrons and ions in that case? Are the input files only giving this option to the electrons?
4. We should add the assert that you mentioned for alpha=pi/2. Also, since the validity of using this outside of the training domain is unknown, I think we should add an fprintf(stdout, with a warning like "sheath_surrogate: outside of region of convergence". In general we stand by what Gkeyll produces, meaning we say its results can be trusted (at least as far as we know). But this option is adding the ability to produce results whose validity is knowningly unknown.
5. Is the gkylcas PR up to date?