Added check of gridding for RZ spectral solver

Python/pywarpx/_libwarpx.py

@@ -248,6 +248,8 @@ def initialize(argv=None):
         argv = sys.argv
     amrex_init(argv)
     libwarpx.warpx_ConvertLabParamsToBoost()
+    if geometry_dim == 'rz':
+        libwarpx.warpx_CheckGriddingForRZSpectral()
     libwarpx.warpx_init()
 
 

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();

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);

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);
 

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);
+    }
+
+    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 two 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);
+
+    // Make 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);
+
+    // Make sure that the max grid size (of the finest 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!")

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();