[PR] FLR effects with Padé approximants

core/zero/cudss_ops.cu

@@ -1,6 +1,7 @@
 #ifdef GKYL_HAVE_CUDSS
 
 #include <cudss.h>
+#include <cusparse.h>
 
 extern "C" {
 #include <gkyl_alloc.h>
@@ -263,6 +264,46 @@ gkyl_culinsolver_sync(struct gkyl_culinsolver_prob *prob)
   cudaStreamSynchronize(prob->stream);
 }
 
+void
+gkyl_culinsolver_mat_vec(struct gkyl_culinsolver_prob *prob, const double *x, double *y)
+{
+  cusparseHandle_t handle;
+  cusparseCreate(&handle);
+  cusparseSetStream(handle, prob->stream);
+
+  double alpha = 1.0, beta = 0.0;
+  void *buf = NULL;
+  size_t buf_sz = 0;
+  for (int k=0; k<prob->nprob; k++) {
+    cusparseSpMatDescr_t matA;
+    cusparseCreateCsr(&matA, prob->mrow, prob->ncol, prob->nnz,
+      prob->csr_rowptr_cu, prob->csr_colind_cu, prob->csr_val_cu+k*prob->nnz,
+      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_64F);
+
+    cusparseDnVecDescr_t vecX, vecY;
+    cusparseCreateDnVec(&vecX, prob->ncol, (void*) &x[k*prob->ncol], CUDA_R_64F);
+    cusparseCreateDnVec(&vecY, prob->mrow, &y[k*prob->mrow], CUDA_R_64F);
+
+    if (k == 0) {
+      // All problems have the same sparsity pattern, so one buffer suffices.
+      cusparseSpMV_bufferSize(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, vecX, &beta, vecY,
+        CUDA_R_64F, CUSPARSE_SPMV_ALG_DEFAULT, &buf_sz);
+      buf = buf_sz > 0? gkyl_cu_malloc(buf_sz) : NULL;
+    }
+
+    cusparseSpMV(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, vecX, &beta, vecY,
+      CUDA_R_64F, CUSPARSE_SPMV_ALG_DEFAULT, buf);
+
+    cusparseDestroyDnVec(vecX);
+    cusparseDestroyDnVec(vecY);
+    cusparseDestroySpMat(matA);
+  }
+  cudaStreamSynchronize(prob->stream);
+
+  if (buf) gkyl_cu_free(buf);
+  cusparseDestroy(handle);
+}
+
 void
 gkyl_culinsolver_finish_host(struct gkyl_culinsolver_prob *prob)
 {

core/zero/cusolver_ops.cu

@@ -425,6 +425,47 @@ gkyl_culinsolver_solve(struct gkyl_culinsolver_prob *prob)
     cusolverRfBatchSolve(prob->cusolverRfH, prob->d_P, prob->d_Q, 1, prob->d_T, prob->mrow, prob->rhspointers_cu, prob->mrow);
 }
 
+void
+gkyl_culinsolver_sync(struct gkyl_culinsolver_prob *prob)
+{
+  cudaStreamSynchronize(prob->stream);
+}
+
+void
+gkyl_culinsolver_mat_vec(struct gkyl_culinsolver_prob *prob, const double *x, double *y)
+{
+  double alpha = 1.0, beta = 0.0;
+  void *buf = NULL;
+  size_t buf_sz = 0;
+  for (int k=0; k<prob->nprob; k++) {
+    cusparseSpMatDescr_t matA;
+    cusparseCreateCsr(&matA, prob->mrow, prob->ncol, prob->nnz,
+      prob->csrrowptrA_cu, prob->csrcolindA_cu, prob->csrvalA_cu+k*prob->nnz,
+      CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_64F);
+
+    cusparseDnVecDescr_t vecX, vecY;
+    cusparseCreateDnVec(&vecX, prob->ncol, (void*) &x[k*prob->ncol], CUDA_R_64F);
+    cusparseCreateDnVec(&vecY, prob->mrow, &y[k*prob->mrow], CUDA_R_64F);
+
+    if (k == 0) {
+      // All problems have the same sparsity pattern, so one buffer suffices.
+      cusparseSpMV_bufferSize(prob->cusparseH, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, vecX, &beta, vecY,
+        CUDA_R_64F, CUSPARSE_SPMV_ALG_DEFAULT, &buf_sz);
+      buf = buf_sz > 0? gkyl_cu_malloc(buf_sz) : NULL;
+    }
+
+    cusparseSpMV(prob->cusparseH, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, vecX, &beta, vecY,
+      CUDA_R_64F, CUSPARSE_SPMV_ALG_DEFAULT, buf);
+
+    cusparseDestroyDnVec(vecX);
+    cusparseDestroyDnVec(vecY);
+    cusparseDestroySpMat(matA);
+  }
+  cudaStreamSynchronize(prob->stream);
+
+  if (buf) gkyl_cu_free(buf);
+}
+
 void
 gkyl_culinsolver_finish_host(struct gkyl_culinsolver_prob *prob)
 {

core/zero/gkyl_cudss_ops.h

@@ -66,6 +66,16 @@ void gkyl_culinsolver_brhs_from_triples(struct gkyl_culinsolver_prob *prob, gkyl
  */
 void gkyl_culinsolver_solve(struct gkyl_culinsolver_prob *prob);
 
+/**
+ * Compute the matrix-vector products y_i = A_i*x_i (on the device) for each
+ * of the nprob problem matrices A_i. Uses cuSPARSE on the CSR arrays.
+ *
+ * @param prob cuDSS struct holding the assembled A matrices.
+ * @param x Input vector on the device (length nprob*ncol).
+ * @param y Output vector on the device (length nprob*mrow).
+ */
+void gkyl_culinsolver_mat_vec(struct gkyl_culinsolver_prob *prob, const double *x, double *y);
+
 /**
  * Copy solution back to host
  *

core/zero/gkyl_cusolver_ops.h

@@ -51,6 +51,23 @@ void gkyl_culinsolver_brhs_from_triples(struct gkyl_culinsolver_prob *prob, gkyl
  */
 void gkyl_culinsolver_solve(struct gkyl_culinsolver_prob *prob);
 
+/**
+ * Synchronize the stream the solver ran on.
+ *
+ * @param prob cuSolver struct holding arrays used in problem.
+ */
+void gkyl_culinsolver_sync(struct gkyl_culinsolver_prob *prob);
+
+/**
+ * Compute the matrix-vector products y_i = A_i*x_i (on the device) for each
+ * of the nprob problem matrices A_i. Uses cuSPARSE on the CSR arrays.
+ *
+ * @param prob cuSolver struct holding the assembled A matrices.
+ * @param x Input vector on the device (length nprob*ncol).
+ * @param y Output vector on the device (length nprob*mrow).
+ */
+void gkyl_culinsolver_mat_vec(struct gkyl_culinsolver_prob *prob, const double *x, double *y);
+
 /**
  * Copy solution back to host
  *

core/zero/gkyl_eqn_type.h

@@ -168,7 +168,8 @@ enum gkyl_vel_edge {
 // Identifiers for FLR models (in gyrokinetics).
 enum gkyl_gk_flr_type {
   GKYL_GK_FLR_NONE = 0, // No FLR effects.
-  GKYL_GK_FLR_PADE_CONST, // Pade-based approx. w/ const. rho_ts=sqrt(Tperp_s/m_s)
+  GKYL_GK_FLR_PADE_CONST, // Pade-based approx. w/ const. rho_ts; at the field level, uses a single reference gyroradius in the operator retrieving phi.
+  GKYL_GK_FLR_PADE_CONST_SUM, // Pade-based approx. w/ const. rho_ts; at the field level, uses the polarization-weighted average of the species gyroradii.
 };
 
 // Gyrokinetic anomaous diffusion models.

core/zero/gkyl_superlu_ops.h

@@ -65,6 +65,16 @@ void gkyl_superlu_brhs_from_array(struct gkyl_superlu_prob *prob, const double *
  */
 void gkyl_superlu_solve(struct gkyl_superlu_prob *prob);
 
+/**
+ * Compute the matrix-vector products y_i = A_i*x_i for each of the nprob
+ * problem matrices A_i. Does not require the A_i to be factorized.
+ *
+ * @param prob SuperLu struct holding the assembled A matrices.
+ * @param x Input vector (length nprob*ncol).
+ * @param y Output vector (length nprob*mrow).
+ */
+void gkyl_superlu_mat_vec(struct gkyl_superlu_prob *prob, const double *x, double *y);
+
 /**
  * Update the values of the left-hand-side SuperLU matrix A in Ax=B problem from a list of
  * triples. Assumes the sparsity pattern didn't change.

core/zero/superlu_ops.c

@@ -319,6 +319,14 @@ gkyl_superlu_solve(struct gkyl_superlu_prob *prob)
   }
 }
 
+void
+gkyl_superlu_mat_vec(struct gkyl_superlu_prob *prob, const double *x, double *y)
+{
+  char trans[2] = "N";
+  for (size_t k=0; k<prob->nprob; k++)
+    sp_dgemv(trans, 1.0, prob->A[k], (double *) &x[k*prob->ncol], 1, 0.0, &y[k*prob->mrow], 1);
+}
+
 void
 gkyl_superlu_amat_update_from_triples(struct gkyl_superlu_prob *prob, struct gkyl_mat_triples **tri)
 {

gyrokinetic/apps/gk_field_1x.c

@@ -17,6 +17,9 @@ gk_field_rhs_phi_1x(struct gkyl_gyrokinetic_app *app, struct gk_field *field)
 {
   // Solve the Poisson equation in 1x with the parallel FEM projection.
   gk_field_fem_projection_par(app, field, field->rho_c, field->phi_smooth);
+
+  // Finish the Poisson solve with FLR effects.
+  field->invert_flr(app, field, field->phi_smooth);
 }
 
 static void

gyrokinetic/apps/gk_field_2x3x.c

@@ -532,33 +532,32 @@ gk_field_fem_new_2x3x(struct gkyl_gyrokinetic_app *app, struct gk_field *f)
   }
 
   // Translate input file BCs into Poisson BCs.
-  struct gkyl_poisson_bc poisson_bcs = { };
   for (int d=0; d<app->cdim-1; d++) {
     if (bc_is_np[d]) {
       struct gkyl_gyrokinetic_bc *bc_lo = gk_fetch_bc_with_dir_edge(f->info.poisson_bcs, 2*app->cdim, d, GKYL_LOWER_EDGE);
       if (bc_lo != 0) {
-        poisson_bcs.lo_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(bc_lo->type);
+        f->poisson_bcs.lo_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(bc_lo->type);
         for (int i=0; i<3; i++) {
-          poisson_bcs.lo_value[d].v[i] = bc_lo->value[i];
+          f->poisson_bcs.lo_value[d].v[i] = bc_lo->value[i];
         }
       }
 
       struct gkyl_gyrokinetic_bc *bc_up = gk_fetch_bc_with_dir_edge(f->info.poisson_bcs, 2*app->cdim, d, GKYL_UPPER_EDGE);
       if (bc_up != 0) {
-        poisson_bcs.up_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(bc_up->type);
+        f->poisson_bcs.up_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(bc_up->type);
         for (int i=0; i<3; i++) {
-          poisson_bcs.up_value[d].v[i] = bc_up->value[i];
+          f->poisson_bcs.up_value[d].v[i] = bc_up->value[i];
         }
       }
     } else {
-      poisson_bcs.lo_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(GKYL_BC_GK_FIELD_PERIODIC);
-      poisson_bcs.up_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(GKYL_BC_GK_FIELD_PERIODIC);
+      f->poisson_bcs.lo_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(GKYL_BC_GK_FIELD_PERIODIC);
+      f->poisson_bcs.up_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(GKYL_BC_GK_FIELD_PERIODIC);
     }
   }
 
   // Initialize the Poisson solver.
   f->fem_poisson_perp = gkyl_fem_poisson_perp_new(&app->local, &app->grid, app->basis,
-    &poisson_bcs, f->info.bias_line_list, f->epsilon, NULL, app->use_gpu);
+    &f->poisson_bcs, f->info.bias_line_list, f->epsilon, NULL, app->use_gpu);
 
   f->phi_bc = 0;
   f->is_dirichletvar = false;
@@ -642,11 +641,11 @@ gk_field_fem_new_2x3x(struct gkyl_gyrokinetic_app *app, struct gk_field *f)
 
   if (app->gk_geom->has_LCFS) {
     // Updaters to enforce twist-and-shift and sheath BCs.
-    gk_field_2x3x_add_IWL_updaters(app, f, &poisson_bcs);
+    gk_field_2x3x_add_IWL_updaters(app, f, &f->poisson_bcs);
   }
   else if (f->bc_par_phi == GKYL_BC_GK_FIELD_TWISTSHIFT) {
     // Updaters to enforce twist-and-shift BCs.
-    gk_field_2x3x_add_TS_updaters(app, f, &poisson_bcs);
+    gk_field_2x3x_add_TS_updaters(app, f, &f->poisson_bcs);
   }
 
   // Set the pointer to the function that computes phi.

gyrokinetic/apps/gk_field_flr.c

@@ -12,59 +12,56 @@ gk_field_flr_new(struct gkyl_gyrokinetic_app *app, struct gk_field *f)
   assert(app->cdim > 1);
   f->invert_flr = gk_field_invert_flr;
 
-  double flr_weight = 0.0;
-  for (int i = 0; i < app->num_species; ++i){
-    struct gk_species *s = &app->species[i];
-    double gyroradius_bmag = s->info.flr.bmag ? s->info.flr.bmag : app->bmag_ref;
-    flr_weight += s->info.flr.Tperp * s->info.mass / (pow(s->info.charge * gyroradius_bmag, 2.0));
+  // Reference (squared) gyroradius in the operator A = 1 - rho^2*nabla_perp^2
+  // used to retrieve phi from the modified potential Phi_0 (step 4 of the
+  // algorithm in flr_effects.tex).
+  double rhoSq_ref = 0.0;
+  if (f->info.flr.type == GKYL_GK_FLR_PADE_CONST) {
+    // Single reference species (e.g. the main ion).
+    assert(f->info.flr.gyroradius > 0.0);
+    rhoSq_ref = pow(f->info.flr.gyroradius, 2.0);
   }
-  // Initialize the weight in the Laplacian operator.
-  f->flr_rhoSq_sum = mkarr(app->use_gpu, (2 * (app->cdim - 1) - 1) * app->basis.num_basis, app->local_ext.volume);
-  gkyl_array_set_offset(f->flr_rhoSq_sum, flr_weight, app->gk_geom->geo_int.gxxj, 0 * app->basis.num_basis);
-  if (app->cdim > 2) {
-    gkyl_array_set_offset(f->flr_rhoSq_sum, flr_weight, app->gk_geom->geo_int.gxyj, 1 * app->basis.num_basis);
-    gkyl_array_set_offset(f->flr_rhoSq_sum, flr_weight, app->gk_geom->geo_int.gyyj, 2 * app->basis.num_basis);
+  else {
+    // polarization-weighted average of the species gyroradii,
+    //   rho^2 = sum_s eps_s0*rho_s0^2 / sum_s eps_s0,  eps_s0 = n_s0*m_s/B^2,
+    double polarization_bmag = f->info.polarization_bmag ? f->info.polarization_bmag : app->bmag_ref;
+    double eps_sum = 0.0;
+    for (int i = 0; i < app->num_species; ++i) {
+      struct gk_species *s = &app->species[i];
+      double gyroradius_bmag = s->info.flr.bmag ? s->info.flr.bmag : app->bmag_ref;
+      double rhoSq_s = s->info.flr.Tperp * s->info.mass / pow(s->info.charge * gyroradius_bmag, 2.0);
+      double eps_s0 = s->info.polarization_density * s->info.mass / pow(polarization_bmag, 2.0);
+      rhoSq_ref += eps_s0 * rhoSq_s;
+      eps_sum += eps_s0;
+    }
+    rhoSq_ref /= eps_sum;
   }
-  // Initialize the factor multiplying the field in the FLR operator.
-  f->flr_kSq = mkarr(app->use_gpu, app->basis.num_basis, app->local_ext.volume);
-  gkyl_array_shiftc(f->flr_kSq, -pow(sqrt(2.0), app->cdim), 0); // Sets kSq=-1.
-
-  // If domain is not periodic use Dirichlet BCs.
-  struct gkyl_poisson_bc flr_bc = {};
-
-  bool bc_is_np[GKYL_MAX_CDIM]; // Is the BC in this direction non-periodic?
-  for (int d = 0; d < app->cdim; ++d) {
-    bc_is_np[d] = true;
+  // The apply passes boundary values through Dirichlet rows, so spatially
+  // varying Dirichlet BCs are not supported with FLR effects.
+  for (int d = 0; d < app->cdim - 1; d++) {
+    assert(f->poisson_bcs.lo_type[d] != GKYL_POISSON_DIRICHLET_VARYING);
+    assert(f->poisson_bcs.up_type[d] != GKYL_POISSON_DIRICHLET_VARYING);
   }
-  for (int d = 0; d < app->num_periodic_dir; ++d) {
-    bc_is_np[app->periodic_dirs[d]] = false;
+
+  // Weight in the perpendicular Laplacian of A = 1 - rho^2*nabla_perp^2.
+  f->flr_rhoSq = mkarr(app->use_gpu, (2*(app->cdim/3)+1)*app->basis.num_basis, app->local_ext.volume);
+  struct gkyl_array *Jgij[3] = {app->gk_geom->geo_int.gxxj, app->gk_geom->geo_int.gxyj, app->gk_geom->geo_int.gyyj};
+  for (int i=0; i<app->cdim-2/app->cdim; i++) {
+    gkyl_array_set_offset(f->flr_rhoSq, rhoSq_ref, Jgij[i], i*app->basis.num_basis);
   }
+  f->flr_kSq = mkarr(app->use_gpu, app->basis.num_basis, app->local_ext.volume);
+  gkyl_array_shiftc(f->flr_kSq, -pow(sqrt(2.0),app->cdim), 0); // Sets kSq=-1.
 
-  for (int d = 0; d < app->cdim - 1; d++) {
-    if (bc_is_np[d]) {
-      struct gkyl_gyrokinetic_bc *bc_lo = gk_fetch_bc_with_dir_edge(f->info.poisson_bcs, 2 * app->cdim, d, GKYL_LOWER_EDGE);
-      if (bc_lo != 0) {
-        flr_bc.lo_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(GKYL_BC_GK_FIELD_DIRICHLET_VARYING);
-      }
+  f->flr_op = gkyl_fem_poisson_perp_new(&app->local, &app->grid, app->basis, &f->poisson_bcs, f->info.bias_line_list, f->flr_rhoSq, f->flr_kSq, app->use_gpu);
 
-      struct gkyl_gyrokinetic_bc *bc_up = gk_fetch_bc_with_dir_edge(f->info.poisson_bcs, 2 * app->cdim, d, GKYL_UPPER_EDGE);
-      if (bc_up != 0) {
-        flr_bc.up_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(GKYL_BC_GK_FIELD_DIRICHLET_VARYING);
-      }
-    } else {
-      flr_bc.lo_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(GKYL_BC_GK_FIELD_PERIODIC);
-      flr_bc.up_type[d] = gkyl_gyrokinetic_translate_poisson_bc_type(GKYL_BC_GK_FIELD_PERIODIC);
-    }
-  }
-  // Deflated Poisson solve is performed on range assuming decomposition is *only* in z.
-  f->flr_op = gkyl_deflated_fem_poisson_new(app->grid, app->basis_on_dev, app->basis,
-                                            app->local, app->local, f->flr_rhoSq_sum, f->flr_kSq, flr_bc, NULL, app->use_gpu);
 }
 
 void
 gk_field_invert_flr(gkyl_gyrokinetic_app *app, struct gk_field *field, struct gkyl_array *phi)
 {
-  gkyl_deflated_fem_poisson_advance(field->flr_op, phi, phi, phi);
+  // Retrieve phi from the potential obtained from the FLR charge.
+  // phi = (1 - rho^2*nabla_perp^2) Phi_0
+  gkyl_fem_poisson_perp_lhs_apply(field->flr_op, phi, phi);
 }
 
 void
@@ -75,7 +72,7 @@ gk_field_invert_flr_none(gkyl_gyrokinetic_app *app, struct gk_field *field, stru
 void
 gk_field_flr_release(const struct gkyl_gyrokinetic_app *app, struct gk_field *f)
 {
-  gkyl_array_release(f->flr_rhoSq_sum);
+  gkyl_array_release(f->flr_rhoSq);
   gkyl_array_release(f->flr_kSq);
-  gkyl_deflated_fem_poisson_release(f->flr_op);
+  gkyl_fem_poisson_perp_release(f->flr_op);
 }