5 files changed, 94 insertions, 0 deletions
diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp
index e87a20a7a..495e07e54 100644
--- a/Source/Python/WarpXWrappers.cpp
+++ b/Source/Python/WarpXWrappers.cpp
@@ -213,6 +213,11 @@ extern "C"
       ConvertLabParamsToBoost();
     }
 
+    void warpx_CheckGriddingForRZSpectral()
+    {
+      CheckGriddingForRZSpectral();
+    }
+
     amrex::Real warpx_getProbLo(int dir)
     {
       WarpX& warpx = WarpX::GetInstance();
diff --git a/Source/Python/WarpXWrappers.h b/Source/Python/WarpXWrappers.h
index 4a649a564..bf1b305f4 100644
--- a/Source/Python/WarpXWrappers.h
+++ b/Source/Python/WarpXWrappers.h
@@ -78,6 +78,8 @@ extern "C" {
 
     void warpx_ConvertLabParamsToBoost();
 
+    void warpx_CheckGriddingForRZSpectral();
+
     amrex::Real warpx_getProbLo(int dir);
 
     amrex::Real warpx_getProbHi(int dir);
diff --git a/Source/Utils/WarpXUtil.H b/Source/Utils/WarpXUtil.H
index bb0a620fa..d3682be2f 100644
--- a/Source/Utils/WarpXUtil.H
+++ b/Source/Utils/WarpXUtil.H
@@ -25,6 +25,11 @@ void ReadBoostedFrameParameters(amrex::Real& gamma_boost, amrex::Real& beta_boos
 
 void ConvertLabParamsToBoost();
 
+/**
+ * \brief Ensures that the blocks are setup correctly for the RZ spectral solver
+ */
+void CheckGriddingForRZSpectral();
+
 void NullifyMF(amrex::MultiFab& mf, int lev, amrex::Real zmin,
                amrex::Real zmax);
 
diff --git a/Source/Utils/WarpXUtil.cpp b/Source/Utils/WarpXUtil.cpp
index 952982c47..c7949925c 100644
--- a/Source/Utils/WarpXUtil.cpp
+++ b/Source/Utils/WarpXUtil.cpp
@@ -213,6 +213,86 @@ WarpXParser makeParser (std::string const& parse_function, std::vector<std::stri
     return parser;
 }
 
+/**
+ * \brief Ensures that the blocks are setup correctly for the RZ spectral solver
+ * When using the RZ spectral solver, the Hankel transform cannot be
+ * divided among multiple blocks. Each block must extend over the
+ * entire radial extent.
+ * The grid can be divided up along z, but the number of blocks
+ * must be >= the number of processors.
+ */
+void CheckGriddingForRZSpectral ()
+{
+#if (defined WARPX_DIM_RZ) && (defined WARPX_USE_PSATD)
+
+    int max_level;
+    Vector<int> n_cell(AMREX_SPACEDIM, -1);
+
+    ParmParse pp_amr("amr");
+
+    pp_amr.get("max_level",max_level);
+    pp_amr.getarr("n_cell",n_cell,0,AMREX_SPACEDIM);
+
+    Vector<int> blocking_factor_x(max_level+1);
+    Vector<int> max_grid_size_x(max_level+1);
+
+    // Set the radial block size to be equal to the radial grid size.
+    blocking_factor_x[0] = n_cell[0];
+    max_grid_size_x[0] = n_cell[0];
+
+    for (int lev=1 ; lev <= max_level ; lev++) {
+        // For this to be correct, this needs to read in any user specified refinement ratios.
+        // But since that is messy and unlikely to be needed anytime soon, the ratio is
+        // fixed to 2 which will be the most likely value.
+        blocking_factor_x[lev] = blocking_factor_x[lev-1]*2; // refRatio(lev-1);
+        max_grid_size_x[lev] = max_grid_size_x[lev-1]*2; // refRatio(lev-1);
+    }
+
+    // Note that any user input values for these parameters are discarded.
+    pp_amr.addarr("blocking_factor_x", blocking_factor_x);
+    pp_amr.addarr("max_grid_size_x", max_grid_size_x);
+
+    // Adjust the longitudinal block sizes, making sure that there are
+    // more blocks than processors.
+    // The factor of 8 is there to make some room for higher order
+    // shape factors and filtering.
+    int nprocs = ParallelDescriptor::NProcs();
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(n_cell[1] >= 8*nprocs,
+                                     "With RZ spectral, there must be at least eight z-cells per processor so that there can be at least one block per processor.");
+
+    // Get the longitudinal blocking factor in case it was set by the user.
+    // If not set, use the default value of 8.
+    Vector<int> bf;
+    pp_amr.queryarr("blocking_factor",bf);
+    pp_amr.queryarr("blocking_factor_y",bf);
+    bf.resize(std::max(static_cast<int>(bf.size()),1),8);
+
+    // Modify the default or any user input, making sure that the blocking factor
+    // is small enough so that there will be at least as many blocks as there are
+    // processors. Because of the ASSERT above, bf will never be less than 8.
+    while (n_cell[1] < nprocs*bf[0]) {
+        bf[0] /= 2;
+    }
+    pp_amr.addarr("blocking_factor_y", bf);
+
+    // Get the longitudinal max grid size in case it was set by the user.
+    // If not set, use the default value of 128.
+    Vector<int> mg;
+    pp_amr.queryarr("max_grid_size",mg);
+    pp_amr.queryarr("max_grid_size_y",mg);
+    mg.resize(std::max(static_cast<int>(mg.size()),1),128);
+
+    // Modify the default or any user input, making sure that the max grid size
+    // (of the coarsest level) is small enough so that there will be at least
+    // as many blocks as there are processors.
+    while (n_cell[1] < nprocs*mg[0]) {
+        mg[0] /= 2;
+    }
+    pp_amr.addarr("max_grid_size_y", mg);
+
+#endif
+}
+
 namespace WarpXUtilMsg{
 
 void AlwaysAssert(bool is_expression_true, const std::string& msg = "ERROR!")
diff --git a/Source/main.cpp b/Source/main.cpp
index 56e84dd9a..ee0db6960 100644
--- a/Source/main.cpp
+++ b/Source/main.cpp
@@ -36,6 +36,8 @@ int main(int argc, char* argv[])
 
     ConvertLabParamsToBoost();
 
+    CheckGriddingForRZSpectral();
+
     WARPX_PROFILE_VAR("main()", pmain);
 
     const Real strt_total = amrex::second();