Coupling with an AI model for sheath BC in GK simulations

gyrokinetic/apps/gk_species.c

@@ -12,6 +12,7 @@
 #include <gkyl_dg_interpolate.h>
 #include <gkyl_translate_dim.h>
 #include <gkyl_proj_on_basis.h>
+#include <gkyl_gk_maxwellian_moments.h>
 
 #include <assert.h>
 #include <time.h>
@@ -196,16 +197,40 @@ gk_species_apply_bc_dynamic(gkyl_gyrokinetic_app *app, const struct gk_species *
   int num_periodic_dir = species->num_periodic_dir, cdim = app->cdim;
   gkyl_comm_array_per_sync(species->comm, &species->local, &species->local_ext,
     num_periodic_dir, species->periodic_dirs, f); 
+
+  if (species->alloc_srg_aux_var) {
+    // Advance the Maxwellian moment computation to get density and temperature for surrogate BCs.
+    // We can later optimize this by calling it for upper and lower skin cells.
+    gkyl_gk_maxwellian_moments_advance(species->maxwell_mom, &species->local, &app->local, f, species->maxmom);
+
+    // Overwrite the density with the polarization density
+    // gkyl_array_set_range(species->dens_sheath, 1.0, species->maxmom, &app->local);
+    gkyl_array_clear(species->dens_sheath, 0.0);
+    gkyl_array_shiftc(species->dens_sheath, species->info.polarization_density * pow(sqrt(2.0),app->cdim), 0);
+
+    int num_conf_basis = app->basis.num_basis;
+    gkyl_array_set_offset_range(species->temp_sheath, species->info.mass, species->maxmom, 2*num_conf_basis, &app->local);
+  }
 
   for (int d=0; d<cdim; ++d) {
     if (species->bc_is_np[d]) {
 
       switch (species->lower_bc[d].type) {
         case GKYL_BC_GK_SPECIES_SHEATH:
+          gkyl_bc_sheath_gyrokinetic_update_vcut_fact_surrogate(species->bc_sheath_lo, 
+            app->field->phi_smooth, app->field->phi_wall_lo, species->dens_sheath, species->temp_sheath, 
+            app->gk_geom->geo_int.bmag, 
+            app->gk_geom->geo_int.bimpactangle, 
+            &app->local);
           gkyl_bc_sheath_gyrokinetic_advance(species->bc_sheath_lo, app->field->phi_smooth, 
             app->field->phi_wall_lo, f, &app->local);
           break;
         case GKYL_BC_GK_SPECIES_IWL:
+          gkyl_bc_sheath_gyrokinetic_update_vcut_fact_surrogate(species->bc_sheath_lo, 
+            app->field->phi_smooth, app->field->phi_wall_lo, species->dens_sheath, species->temp_sheath, 
+            app->gk_geom->geo_int.bmag, 
+            app->gk_geom->geo_int.bimpactangle, 
+            &app->local);
           gkyl_bc_sheath_gyrokinetic_advance(species->bc_sheath_lo, app->field->phi_smooth, 
             app->field->phi_wall_lo, f, &app->local);
           if (cdim == 3) {
@@ -228,10 +253,20 @@ gk_species_apply_bc_dynamic(gkyl_gyrokinetic_app *app, const struct gk_species *
 
       switch (species->upper_bc[d].type) {
         case GKYL_BC_GK_SPECIES_SHEATH:
+          gkyl_bc_sheath_gyrokinetic_update_vcut_fact_surrogate(species->bc_sheath_up, 
+            app->field->phi_smooth, app->field->phi_wall_up, species->dens_sheath, species->temp_sheath, 
+            app->gk_geom->geo_int.bmag, 
+            app->gk_geom->geo_int.bimpactangle, 
+            &app->local);
           gkyl_bc_sheath_gyrokinetic_advance(species->bc_sheath_up, app->field->phi_smooth, 
             app->field->phi_wall_up, f, &app->local);
           break;
         case GKYL_BC_GK_SPECIES_IWL:
+          gkyl_bc_sheath_gyrokinetic_update_vcut_fact_surrogate(species->bc_sheath_up, 
+            app->field->phi_smooth, app->field->phi_wall_up, species->dens_sheath, species->temp_sheath, 
+            app->gk_geom->geo_int.bmag, 
+            app->gk_geom->geo_int.bimpactangle, 
+            &app->local);
           gkyl_bc_sheath_gyrokinetic_advance(species->bc_sheath_up, app->field->phi_smooth, 
             app->field->phi_wall_up, f, &app->local);
           if (cdim == 3) {
@@ -833,12 +868,12 @@ gk_species_init_dynamic(struct gkyl_gk *gk_app_inp, struct gkyl_gyrokinetic_app
     if (gks->lower_bc[d].type == GKYL_BC_GK_SPECIES_SHEATH) {
       gks->bc_sheath_lo = gkyl_bc_sheath_gyrokinetic_new(d, GKYL_LOWER_EDGE, gks->basis_on_dev, 
         &gks->lower_skin[d], &gks->lower_ghost[d], gks->vel_map,
-        cdim, 2.0*(gks->info.charge/gks->info.mass), app->use_gpu);
+        cdim, 2.0*(gks->info.charge/gks->info.mass), gks->lower_bc[d].use_sheath_surrogate, app->use_gpu);
     }
     else if (gks->lower_bc[d].type == GKYL_BC_GK_SPECIES_IWL) {
       gks->bc_sheath_lo = gkyl_bc_sheath_gyrokinetic_new(d, GKYL_LOWER_EDGE, gks->basis_on_dev, 
         &gks->lower_skin_par_sol, &gks->lower_ghost_par_sol, gks->vel_map,
-        cdim, 2.0*(gks->info.charge/gks->info.mass), app->use_gpu);
+        cdim, 2.0*(gks->info.charge/gks->info.mass), gks->lower_bc[d].use_sheath_surrogate, app->use_gpu);
 
       if (cdim == 3) {
         // For 3x2v we need a twistshift BC in the core.
@@ -893,12 +928,12 @@ gk_species_init_dynamic(struct gkyl_gk *gk_app_inp, struct gkyl_gyrokinetic_app
     if (gks->upper_bc[d].type == GKYL_BC_GK_SPECIES_SHEATH) {
       gks->bc_sheath_up = gkyl_bc_sheath_gyrokinetic_new(d, GKYL_UPPER_EDGE, gks->basis_on_dev, 
         &gks->upper_skin[d], &gks->upper_ghost[d], gks->vel_map,
-        cdim, 2.0*(gks->info.charge/gks->info.mass), app->use_gpu);
+        cdim, 2.0*(gks->info.charge/gks->info.mass), gks->upper_bc[d].use_sheath_surrogate, app->use_gpu);
     }
     else if (gks->upper_bc[d].type == GKYL_BC_GK_SPECIES_IWL) {
       gks->bc_sheath_up = gkyl_bc_sheath_gyrokinetic_new(d, GKYL_UPPER_EDGE, gks->basis_on_dev, 
         &gks->upper_skin_par_sol, &gks->upper_ghost_par_sol, gks->vel_map,
-        cdim, 2.0*(gks->info.charge/gks->info.mass), app->use_gpu);
+        cdim, 2.0*(gks->info.charge/gks->info.mass), gks->upper_bc[d].use_sheath_surrogate, app->use_gpu);
 
       if (cdim == 3) {
         // For 3x2v we need a twistshift BC in the core.
@@ -950,6 +985,30 @@ gk_species_init_dynamic(struct gkyl_gk *gk_app_inp, struct gkyl_gyrokinetic_app
     }
   }
 
+  gks->alloc_srg_aux_var = false;
+  for (int d=0; d<app->cdim; ++d) {
+    gks->alloc_srg_aux_var = gks->alloc_srg_aux_var || gks->lower_bc[d].use_sheath_surrogate;
+    gks->alloc_srg_aux_var = gks->alloc_srg_aux_var || gks->upper_bc[d].use_sheath_surrogate;    
+  }
+  if (gks->alloc_srg_aux_var) {
+    struct gkyl_gk_maxwellian_moments_inp inp_mom = {
+      .phase_grid = &gks->grid,
+      .conf_basis = &app->basis,
+      .phase_basis = &gks->basis,
+      .conf_range =  &app->local,
+      .conf_range_ext = &app->local_ext,
+      .mass = gks->info.mass, 
+      .gk_geom = app->gk_geom, 
+      .vel_map = gks->vel_map,
+      .divide_jacobgeo = true, 
+      .use_gpu = app->use_gpu,
+    };
+    gks->maxwell_mom = gkyl_gk_maxwellian_moments_inew( &inp_mom );
+    gks->maxmom = mkarr(app->use_gpu, 3*app->basis.num_basis, app->local_ext.volume);
+    gks->dens_sheath = mkarr(app->use_gpu, app->basis.num_basis, app->local_ext.volume);
+    gks->temp_sheath = mkarr(app->use_gpu, app->basis.num_basis, app->local_ext.volume);
+  }
+
   // Set function pointers.
   gks->rhs_func = gk_species_rhs_dynamic;
   gks->rhs_implicit_func = gk_species_rhs_implicit_dynamic;
@@ -1938,5 +1997,9 @@ gk_species_release(const gkyl_gyrokinetic_app* app, const struct gk_species *gks
     gkyl_deflated_fem_poisson_release(gks->flr_op);
   }
 
+  if (gks->alloc_srg_aux_var) {
+    gkyl_gk_maxwellian_moments_release(gks->maxwell_mom);
+  }
+
   gks->release_func(app, gks);
 }

gyrokinetic/apps/gkyl_gyrokinetic.h

@@ -178,6 +178,7 @@ struct gkyl_gyrokinetic_bc {
   struct gkyl_gyrokinetic_emission_inp emission; 
   bool write_diagnostics; // Whether to output diagnostics.
   int bidx; // Block index (for multiblock solver).
+  bool use_sheath_surrogate; // Whether to use surrogate sheath BC to determine vcut (only used if type is SHEATH).
 };
 
 struct gkyl_gyrokinetic_geometry {

gyrokinetic/apps/gkyl_gyrokinetic_priv.h

@@ -946,7 +946,13 @@ struct gk_species {
   struct gkyl_array *f_host; // Host copy for IO and initialization.
 
   struct gkyl_array *gyro_phi; // Gyroaveraged electrostatic potential.
-
+
+  bool alloc_srg_aux_var; // To know whether surrogate auxiliary variables are allocated.
+  struct gkyl_gk_maxwellian_moments *maxwell_mom; // Maxwellian moment array to get temperature and density for surrogate BCs.
+  struct gkyl_array *maxmom; // Array for maxwellian moments.
+  struct gkyl_array *dens_sheath; // Zeroth maxwellian moment, at the sheath used for surrogate BCs.
+  struct gkyl_array *temp_sheath; // Third maxwellian moment times mass, at the sheath used for surrogate BCs.
+
   struct gk_species_moment m0; // Computes charge density.
   struct gk_species_moment integ_moms; // Integrated moments.
   struct gk_species_moment *moms; // Diagnostic moments

gyrokinetic/apps/gyrokinetic.c

@@ -1081,6 +1081,7 @@ gkyl_gyrokinetic_app_write_geometry(gkyl_gyrokinetic_app* app, struct gkyl_gk_ge
   gyrokinetic_app_geometry_copy_and_write(app, app->gk_geom->geo_int.bioverJB, arr_ho3, "bioverJB", mt);
   gyrokinetic_app_geometry_copy_and_write(app, app->gk_geom->geo_int.B3, arr_ho1, "B3", mt);
   gyrokinetic_app_geometry_copy_and_write(app, app->gk_geom->geo_int.dualcurlbhat, arr_ho3, "dualcurlbhat", mt);
+  gyrokinetic_app_geometry_copy_and_write(app, app->gk_geom->geo_int.bimpactangle, arr_ho1, "bimpactangle", mt);
 
 
   // Write surface quantities
@@ -2717,6 +2718,7 @@ gkyl_gyrokinetic_app_read_geometry(gkyl_gyrokinetic_app* app)
   gyrokinetic_app_geometry_read_and_copy(app, app->gk_geom->geo_int.bioverJB         , arr_ho3, "bioverJB");
   gyrokinetic_app_geometry_read_and_copy(app, app->gk_geom->geo_int.B3               , arr_ho1, "B3");
   gyrokinetic_app_geometry_read_and_copy(app, app->gk_geom->geo_int.dualcurlbhat     , arr_ho3, "dualcurlbhat");
+  gyrokinetic_app_geometry_read_and_copy(app, app->gk_geom->geo_int.bimpactangle     , arr_ho1, "bimpactangle");
 
   struct gkyl_array* arr_surf_ho1 = mkarr(false,   app->gk_geom->num_surf_basis, app->local_ext.volume);
   struct gkyl_array* arr_surf_ho2 = mkarr(false, 2*app->gk_geom->num_surf_basis, app->local_ext.volume);