Kernels for AI sheath interface

maxima/g0/gk-sheath/ms-gk-sheath.mac

@@ -8,9 +8,14 @@ load("gk-sheath/sheath-1x1v");
 load("gk-sheath/sheath-1x2v");
 load("gk-sheath/sheath-2x2v");
 load("gk-sheath/sheath-3x2v");
+load("gk-sheath/sheath_surrogate-1x2v");
+load("gk-sheath/sheath_surrogate-2x2v");
+load("gk-sheath/sheath_surrogate-3x2v");
 
 /* ...... USER INPUTS........ */
 
+outDir : "~/max-out"$ /* Output directory */
+
 /* Serendipity basis. */
 minPolyOrder_Ser : 1$
 maxPolyOrder_Ser : 1$
@@ -43,7 +48,7 @@ for bInd : 1 thru length(bName) do (
   minPolyOrderB : minPolyOrder[bInd],
   maxPolyOrderB : maxPolyOrder[bInd],
   for polyOrder : minPolyOrderB thru maxPolyOrderB do (
-    fname : sconcat("~/max-out/bc_sheath_gyrokinetic_",bName[bInd],"_p",polyOrder,".c"),
+    fname : sconcat(outDir,"/bc_sheath_gyrokinetic_",bName[bInd],"_p",polyOrder,".c"),
     fh : openw(fname),
 
     printf(fh, "#include <gkyl_bc_sheath_gyrokinetic_kernels.h> ~%"),
@@ -63,15 +68,45 @@ for bInd : 1 thru length(bName) do (
         )
       )
     ),
+
+    for useSurrogate : 0 thru 1 do (
+      for cdim : minDim[bInd] thru maxDim[bInd] do (
+        print(sconcat("Creating vcut kernels for serendipity basis, surrogate = ", useSurrogate, " for ", cdim, "x2v, p", polyOrder)),
+        if (polyOrder = 1) then (
+          if cdim=1 then (
+            genGkSheathVcutCalcKer1x2v(fh, cdim, 2, bName[bInd], polyOrder, useSurrogate)
+          ) else if cdim=2 then (
+            genGkSheathVcutCalcKer2x2v(fh, cdim, 2, bName[bInd], polyOrder, useSurrogate)
+          ) else if cdim=3 then (
+            genGkSheathVcutCalcKer3x2v(fh, cdim, 2, bName[bInd], polyOrder, useSurrogate)
+          )
+        )
+      )
+    ),
+
+    for cdim : minDim[bInd] thru maxDim[bInd] do (
+      print(sconcat("Creating input builder kernels for serendipity basis for ", cdim, "x2v, p", polyOrder)),
+      if (polyOrder = 1) then (
+        if cdim=1 then (
+          genGkSheathInputKer1x2v(fh, cdim, 2, bName[bInd], polyOrder)
+        ) else if cdim=2 then (
+          genGkSheathInputKer2x2v(fh, cdim, 2, bName[bInd], polyOrder)
+        ) else if cdim=3 then (
+          genGkSheathInputKer3x2v(fh, cdim, 2, bName[bInd], polyOrder)
+        )
+      )
+    ),
     close(fh)
   )
 )$
 
 
 /* Create a header file for sheath kernels. */
-fh : openw("~/max-out/gkyl_bc_sheath_gyrokinetic_kernels.h")$
+fh : openw(sconcat(outDir,"/gkyl_bc_sheath_gyrokinetic_kernels.h"))$
 printf(fh, "#pragma once~%")$
 printf(fh, "~%")$
+printf(fh, "#include <gkyl_const.h>~%")$
+printf(fh, "#include <gkyl_kann_net.h>~%")$
 printf(fh, "#include <gkyl_util.h>~%")$
 printf(fh, "#include <math.h>~%")$
 printf(fh, "~%")$
@@ -83,7 +118,31 @@ printf(fh, "  // from Kroger ~%")$
 printf(fh, "  return (3.*x-x*x*x*(6. + x*x - 2.*x*x*x*x)/5.)/(1.-x*x); ~%")$
 printf(fh, "}~%")$
 printf(fh, "~%")$
-
+
+printf(fh, "// Linear interpolation of the surrogate output vcut[ng] (defined on the mu-grid~%")$
+printf(fh, "// mu_grid[ng], which the surrogate model supplies as the second half of each~%")$
+printf(fh, "// node's output block) onto mu_new[n], normalising mu by mu_ref (= T/B). Results~%")$
+printf(fh, "// are written into out[n]; clamps at the grid boundaries.~%")$
+printf(fh, "GKYL_CU_DH~%")$
+printf(fh, "static inline void~%")$
+printf(fh, "bc_sheath_gyrokinetic_surr_interpf(const float *vcut, const float *mu_grid, int ng, const double *mu_new, int n, double mu_ref, double *out)~%")$
+printf(fh, "{~%")$
+printf(fh, "  for (int i = 0; i < n; i++) {~%")$
+printf(fh, "    double mu = mu_new[i]/mu_ref;~%")$
+printf(fh, "    if (mu <= mu_grid[0])      { out[i] = vcut[0];      continue; }~%")$
+printf(fh, "    if (mu >= mu_grid[ng - 1]) { out[i] = vcut[ng - 1]; continue; }~%")$
+printf(fh, "    // binary search for the bracketing interval~%")$
+printf(fh, "    int lo = 0, hi = ng - 1;~%")$
+printf(fh, "    while (hi - lo > 1) {~%")$
+printf(fh, "      int mid = (lo + hi) >> 1;~%")$
+printf(fh, "      if (mu_grid[mid] <= mu) lo = mid; else hi = mid;~%")$
+printf(fh, "    }~%")$
+printf(fh, "    double t = (mu - mu_grid[lo]) / (mu_grid[hi] - mu_grid[lo]);~%")$
+printf(fh, "    out[i] = vcut[lo] + t * (vcut[hi] - vcut[lo]);~%")$
+printf(fh, "  }~%")$
+printf(fh, "}~%")$
+printf(fh, "~%")$
+
 printf(fh, "EXTERN_C_BEG~%")$
 printf(fh, "~%")$
 
@@ -94,14 +153,35 @@ for bInd : 1 thru length(bName) do (
       for vI : 1 thru length(vDims[cdim]) do (
         for sI : 1 thru length(edge) do (
 
-          printf(fh, "GKYL_CU_DH void bc_sheath_gyrokinetic_reflectedf_~a_~ax~av_~a_p~a(const double *vmap, const double q2Dm, const double *phi, const double *phiWall, const double *f, double *fRefl); ~%", edge[sI], cdim, vDims[cdim][vI], bName[bInd], polyOrder)
+          printf(fh, "GKYL_CU_DH void bc_sheath_gyrokinetic_reflectedf_~a_~ax~av_~a_p~a(const double *vmap, const double *vcutsq, const double *f, double *fRefl); ~%", edge[sI], cdim, vDims[cdim][vI], bName[bInd], polyOrder)
 
         )
       )
     )
   )
 )$
 
+/* Eval vcutconst kernels. */
+for cdim : minDim_Ser thru maxDim_Ser do (
+  for sI : 1 thru length(edge) do (
+    printf(fh, "GKYL_CU_DH void bc_sheath_gyrokinetic_vcutsq_const_~a_~ax2v_ser_p1(const double *phi, const double *phi_wall, double q2Dm, double *vcutSq_n) ; ~%", edge[sI], cdim)
+  )
+)$
+
+/* Eval vcutsq surrogate kernels. */
+for cdim : minDim_Ser thru maxDim_Ser do (
+  for sI : 1 thru length(edge) do (
+    printf(fh, "GKYL_CU_DH void bc_sheath_gyrokinetic_vcutsq_surr_~a_~ax2v_ser_p1(const double *vmap, const float *nn_out, int n_out, const double *temperature, const double *bmag, double *vcutsq_out) ; ~%", edge[sI], cdim)
+  )
+)$
+
+/* input builder kernels. */
+for cdim : minDim_Ser thru maxDim_Ser do (
+  for sI : 1 thru length(edge) do (
+    printf(fh, "GKYL_CU_DH void bc_sheath_gyrokinetic_build_input_~a_~ax2v_ser_p1(const double *phi, const double *phi_wall, const double *density, const double *temperature, const double *bmag, const double *bimpact_angle, int n_inp, float *nn_inp_out) ; ~%", edge[sI], cdim)
+  )
+)$
+
 printf(fh, "~%")$
 printf(fh, "EXTERN_C_END~%")$
 printf(fh, "~%")$

maxima/g0/gk-sheath/sheath-1x1v.mac

@@ -11,7 +11,7 @@ fpprec : 24$
 genGkSheathKer1x1v(fh, cdim, vdim, basisFun, polyOrder) := block(
   [pdim,varsC,bC,varsP,bP,vSub,NP,varsV,vmap_e,vmapSq_e,vmap_prime_e,vparLo_e,vparUp_e,surfPerpVar,
    surfVars,surfVarsC,bSurf,bV,bVpar,bMu,numBsurf,zvSurfVars,nodesMu,numNodesMu,bNMu,f_e,edge,
-   surfPerpVal,sI,phiSheath_e,phiWall_e,deltaPhi_e,deltaPhi_q,vcutSq_q,fSurf_c,fSurf_e,fSurfMu_q,
+   surfPerpVal,sI,phiSheath_e,phiWall_e,deltaPhi_e,deltaPhi_q,vcutsq_q,fSurf_c,fSurf_e,fSurfMu_q,
    fReflSurfMu_e,j,nodeIdx,fSurfMu_c,fSurfMu_e,intLims,fReflSurfMu_c,
    xBar_v,xSqBar_v,fReflSurf_c,fReflSurf_e,fRefl_c], 
 
@@ -54,24 +54,13 @@ genGkSheathKer1x1v(fh, cdim, vdim, basisFun, polyOrder) := block(
   for sI : 1 thru length(edge) do (
 
     printf(fh, "~%"),
-    printf(fh, "GKYL_CU_DH void bc_sheath_gyrokinetic_reflectedf_~a_~ax~av_~a_p~a(const double *vmap, const double q2Dm, const double *phi, const double *phiWall, const double *f, double *fRefl) ~%{ ~%", edge[sI], cdim, vdim, basisFun, polyOrder),
-
-    /* Calculate expansion for deltaPhi = phiSheath - phiWall, evaluated at z=zVal, 
-       where zVal=-1 or 1 depending on if we're on the left or right edge of the global domain, respectively. */
-    phiSheath_e : subst(surfPerpVar=surfPerpVal[sI],doExpand1(phi, bC)),
-    phiWall_e   : subst(surfPerpVar=surfPerpVal[sI],doExpand1(phiWall, bC)),
-    deltaPhi_e  : phiSheath_e - phiWall_e,
-
-    /* Evaluate vcutSqQ = vcut^2. q2Dm = q*2/m. */
-    deltaPhi_q : deltaPhi_e,
-    vcutSq_q   : gcfac(float(fullratsimp(-q2Dm*deltaPhi_q))),
-
+    printf(fh, "GKYL_CU_DH void bc_sheath_gyrokinetic_reflectedf_~a_~ax~av_~a_p~a(const double *vmap, const double *vcutsq, const double *f, double *fRefl) ~%{ ~%", edge[sI], cdim, vdim, basisFun, polyOrder),
+
     /* Evaluate f at the z surface. */
     fSurf_c : calcInnerProdList(surfVars, 1, bSurf, subst(surfPerpVar=surfPerpVal[sI], f_e)),
     fSurf_e : doExpand(fSurf_c, bSurf),
 
     /* Variable declarations/allocations. */
-    printf(fh, "  double vcutSq;~%"),
     printf(fh, "  double fReflSurf[~a] = {0.}; ~%", length(bV)),
     printf(fh, "~%"),
 
@@ -85,21 +74,21 @@ genGkSheathKer1x1v(fh, cdim, vdim, basisFun, polyOrder) := block(
     printf(fh, "  double vparAbsSqUp = vmap[0]>0.? vparUp*vparUp : vparLo*vparLo;~%"),
     printf(fh, "~%"),
 
-    /* Write vcut^2 at boundary. */
-    printf(fh, "  vcutSq = ~a; ~%", vcutSq_q),
+    /* vcutsq is passed as a pre-computed DG field; read the single surface value. */
+    printf(fh, "  double vcutsq_n = vcutsq[0]; ~%"),
     printf(fh, "~%"),
 
     /* If vcut^2 at this node is below all vpar^2 in this cell, BC at this node should be absorbing
        so set coefficients of fRefl to 0 (no reflection from this node). */
-    printf(fh, "  if (vcutSq <= vparAbsSqLo) { // absorb (no reflection) ~%"),
+    printf(fh, "  if (vcutsq_n <= vparAbsSqLo) { // absorb (no reflection) ~%"),
     printf(fh, "~%"),
 
     writeCExprsWithZeros1(fReflSurf, makelist(0.,j,1,length(bV))),
     printf(fh, "~%"),
 
     /* If vcut^2 at this node is above all vpar^2 in this cell, BC at this node should be full reflection. */
     /* So set coefficients of fRefl to coefficients of f. */
-    printf(fh, "  } else if (vcutSq > vparAbsSqUp) { // full reflection ~%"),
+    printf(fh, "  } else if (vcutsq_n > vparAbsSqUp) { // full reflection ~%"),
     printf(fh, "~%"),
 
     /* Project f onto vpar basis (bV). */
@@ -124,7 +113,7 @@ genGkSheathKer1x1v(fh, cdim, vdim, basisFun, polyOrder) := block(
 
     printf(fh, "    if (wv > 0.) {~%"),
     printf(fh, "      // vcut in logical space.~%"),
-    printf(fh, "      double vcut_l = 2.*(sqrt(vcutSq)-wv)/dv;~%"),
+    printf(fh, "      double vcut_l = 2.*(sqrt(vcutsq_n)-wv)/dv;~%"),
     printf(fh, "~%"),
 
     intLims : [[-1,vcut_l]],
@@ -135,7 +124,7 @@ genGkSheathKer1x1v(fh, cdim, vdim, basisFun, polyOrder) := block(
 
     printf(fh, "    } else {~%"),
     printf(fh, "      // vcut in logical space.~%"),
-    printf(fh, "      double vcut_l = 2.*(-sqrt(vcutSq)-wv)/dv;~%"),
+    printf(fh, "      double vcut_l = 2.*(-sqrt(vcutsq_n)-wv)/dv;~%"),
     printf(fh, "~%"),
 
     intLims : [[vcut_l,1]],