Bug fixes and cleanup in load balancing (#2563)

* added helper function to rebuild MultiFabs and iMultiFabs during load balancing and included rebuilding of EB multifabs

* added redistribute call for the particle boundary buffer during load balancing

* consistently use DistribtionMap rather than dmap in ElectrostaticSolver.cpp

* applied suggested changes from code review by Phil Miller

* removed default argument for redistribute in RemakeMultiFab

* removed RemakeMultiFab() as a member of WarpX

* Only remake EB multifabs if they are used

Co-authored-by: Lorenzo Giacomel <47607756+lgiacome@users.noreply.github.com>

* adapted existing particle scraping test (PICMI version) to also cover the redistribution of particle buffers from load balancing

* added redeclaring of m_borrowing

* Move redeclaring of m_borrow inside if statement for ECT solver algorihtm

Co-authored-by: Lorenzo Giacomel <47607756+lgiacome@users.noreply.github.com>

* added calls to MarkCells and ComputeFaceExtensions

* fixed issue causing CI test to fail and copied conditionals from WarpXInitData.cpp to recompute EB quantities

* Guard cells communication for EB data when re-gridding (#105)

* Add 2D circle EB test (#2538)

* Added embedded_circle test

* Add embedded_circle test files

* Removed diag files

* removed PICMI input file

* Update to use default regression analysis

* Added line breaks for spacing

* Added description

* Fixed benchmark file

* Added load balancing to test

* Commented out load_balancing portion of test.
This will be added back in once load balancing is fixed.

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* Added guard cells communication for EB data in regridding

Co-authored-by: Kevin Z. Zhu <86268612+KZhu-ME@users.noreply.github.com>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* moved all EB grid data calculations to a new function InitializeEBGridData() which is now called by both WarpX::InitLevelData and WarpX::RemakeLevel

* Fix typo in doc string.

Co-authored-by: Phil Miller <unmobile+gh@gmail.com>

Co-authored-by: Lorenzo Giacomel <47607756+lgiacome@users.noreply.github.com>
Co-authored-by: Kevin Z. Zhu <86268612+KZhu-ME@users.noreply.github.com>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
Co-authored-by: Phil Miller <unmobile+gh@gmail.com>

1 files changed, 66 insertions, 162 deletions

diff --git a/Source/Parallelization/WarpXRegrid.cpp b/Source/Parallelization/WarpXRegrid.cpp
index 06e8b99ea..c96d5e1c7 100644
--- a/Source/Parallelization/WarpXRegrid.cpp
+++ b/Source/Parallelization/WarpXRegrid.cpp
@@ -11,6 +11,7 @@
 #include "Diagnostics/MultiDiagnostics.H"
 #include "Diagnostics/ReducedDiags/MultiReducedDiags.H"
 #include "Particles/MultiParticleContainer.H"
+#include "Particles/ParticleBoundaryBuffer.H"
 #include "Particles/WarpXParticleContainer.H"
 #include "Utils/WarpXAlgorithmSelection.H"
 #include "Utils/WarpXProfilerWrapper.H"
@@ -131,6 +132,9 @@ WarpX::LoadBalance ()
         mypc->Redistribute();
         mypc->defineAllParticleTiles();
 
+        // redistribute particle boundary buffer
+        m_particle_boundary_buffer->redistribute();
+
         // diagnostics & reduced diagnostics
         // not yet needed:
         //multi_diags->LoadBalance();
@@ -140,6 +144,17 @@ WarpX::LoadBalance ()
 }
 
 
+template <typename MultiFabType> void
+RemakeMultiFab (std::unique_ptr<MultiFabType>& mf, const DistributionMapping& dm,
+                const bool redistribute)
+{
+    if (mf == nullptr) return;
+    const IntVect& ng = mf->nGrowVect();
+    auto pmf = std::make_unique<MultiFabType>(mf->boxArray(), dm, mf->nComp(), ng);
+    if (redistribute) pmf->Redistribute(*mf, 0, 0, mf->nComp(), ng);
+    mf = std::move(pmf);
+}
+
 void
 WarpX::RemakeLevel (int lev, Real /*time*/, const BoxArray& ba, const DistributionMapping& dm)
 {
@@ -147,75 +162,49 @@ WarpX::RemakeLevel (int lev, Real /*time*/, const BoxArray& ba, const Distributi
     {
         if (ParallelDescriptor::NProcs() == 1) return;
 
-#ifdef AMREX_USE_EB
-        m_field_factory[lev] = amrex::makeEBFabFactory(Geom(lev), ba, dm,
-                                                       {1,1,1}, // Not clear how many ghost cells we need yet
-                                                       amrex::EBSupport::full);
-        ComputeEdgeLengths();
-        ComputeFaceAreas();
-        ScaleEdges();
-        ScaleAreas();
-        ComputeDistanceToEB();
-#else
-        m_field_factory[lev] = std::make_unique<FArrayBoxFactory>();
-#endif
-
         // Fine patch
         for (int idim=0; idim < 3; ++idim)
         {
-            {
-                const IntVect& ng = Bfield_fp[lev][idim]->nGrowVect();
-                auto pmf = std::make_unique<MultiFab>(Bfield_fp[lev][idim]->boxArray(),
-                                                                  dm, Bfield_fp[lev][idim]->nComp(), ng);
-                pmf->Redistribute(*Bfield_fp[lev][idim], 0, 0, Bfield_fp[lev][idim]->nComp(), ng);
-                Bfield_fp[lev][idim] = std::move(pmf);
-            }
-            {
-                const IntVect& ng = Efield_fp[lev][idim]->nGrowVect();
-                auto pmf = std::make_unique<MultiFab>(Efield_fp[lev][idim]->boxArray(),
-                                                      dm, Efield_fp[lev][idim]->nComp(), ng);
-                pmf->Redistribute(*Efield_fp[lev][idim], 0, 0, Efield_fp[lev][idim]->nComp(), ng);
-                Efield_fp[lev][idim] = std::move(pmf);
-            }
-            {
-                const IntVect& ng = current_fp[lev][idim]->nGrowVect();
-                auto pmf = std::make_unique<MultiFab>(current_fp[lev][idim]->boxArray(),
-                                                      dm, current_fp[lev][idim]->nComp(), ng);
-                current_fp[lev][idim] = std::move(pmf);
-            }
-            if (current_store[lev][idim])
-            {
-                const IntVect& ng = current_store[lev][idim]->nGrowVect();
-                auto pmf = std::make_unique<MultiFab>(current_store[lev][idim]->boxArray(),
-                                                      dm, current_store[lev][idim]->nComp(), ng);
-                // no need to redistribute
-                current_store[lev][idim] = std::move(pmf);
+            RemakeMultiFab(Bfield_fp[lev][idim], dm, true);
+            RemakeMultiFab(Efield_fp[lev][idim], dm, true);
+            RemakeMultiFab(current_fp[lev][idim], dm, false);
+            RemakeMultiFab(current_store[lev][idim], dm, false);
+
+#ifdef AMREX_USE_EB
+            if (WarpX::maxwell_solver_id == MaxwellSolverAlgo::Yee ||
+                WarpX::maxwell_solver_id == MaxwellSolverAlgo::ECT ||
+                WarpX::maxwell_solver_id == MaxwellSolverAlgo::CKC){
+                RemakeMultiFab(m_edge_lengths[lev][idim], dm, false);
+                RemakeMultiFab(m_face_areas[lev][idim], dm, false);
+                if(WarpX::maxwell_solver_id == MaxwellSolverAlgo::ECT){
+                    RemakeMultiFab(Venl[lev][idim], dm, false);
+                    RemakeMultiFab(m_flag_info_face[lev][idim], dm, false);
+                    RemakeMultiFab(m_flag_ext_face[lev][idim], dm, false);
+                    RemakeMultiFab(m_area_mod[lev][idim], dm, false);
+                    RemakeMultiFab(ECTRhofield[lev][idim], dm, false);
+                    m_borrowing[lev][idim] = std::make_unique<amrex::LayoutData<FaceInfoBox>>(amrex::convert(ba, Bfield_fp[lev][idim]->ixType().toIntVect()), dm);
+                }
             }
+#endif
         }
 
-        if (F_fp[lev] != nullptr) {
-            const IntVect& ng = F_fp[lev]->nGrowVect();
-            auto pmf = std::make_unique<MultiFab>(F_fp[lev]->boxArray(),
-                                                              dm, F_fp[lev]->nComp(), ng);
-            pmf->Redistribute(*F_fp[lev], 0, 0, F_fp[lev]->nComp(), ng);
-            F_fp[lev] = std::move(pmf);
-        }
+        RemakeMultiFab(F_fp[lev], dm, true);
+        RemakeMultiFab(rho_fp[lev], dm, false);
+        // phi_fp should be redistributed since we use the solution from
+        // the last step as the initial guess for the next solve
+        RemakeMultiFab(phi_fp[lev], dm, true);
 
-        if (rho_fp[lev] != nullptr) {
-            const int nc = rho_fp[lev]->nComp();
-            const IntVect& ng = rho_fp[lev]->nGrowVect();
-            auto pmf = std::make_unique<MultiFab>(rho_fp[lev]->boxArray(),
-                                                              dm, nc, ng);
-            rho_fp[lev] = std::move(pmf);
-        }
+#ifdef AMREX_USE_EB
+        RemakeMultiFab(m_distance_to_eb[lev], dm, false);
 
-        if (phi_fp[lev] != nullptr) {
-            const int nc = phi_fp[lev]->nComp();
-            const IntVect& ng = phi_fp[lev]->nGrowVect();
-            auto pmf = std::make_unique<MultiFab>(phi_fp[lev]->boxArray(),
-                                                              dm, nc, ng);
-            phi_fp[lev] = std::move(pmf);
-        }
+        m_field_factory[lev] = amrex::makeEBFabFactory(Geom(lev), ba, dm,
+                                                       {1,1,1}, // Not clear how many ghost cells we need yet
+                                                       amrex::EBSupport::full);
+
+        InitializeEBGridData(lev);
+#else
+        m_field_factory[lev] = std::make_unique<FArrayBoxFactory>();
+#endif
 
 #ifdef WARPX_USE_PSATD
         if (maxwell_solver_id == MaxwellSolverAlgo::PSATD) {
@@ -261,20 +250,8 @@ WarpX::RemakeLevel (int lev, Real /*time*/, const BoxArray& ba, const Distributi
         } else {
             for (int idim=0; idim < 3; ++idim)
             {
-                {
-                    const IntVect& ng = Bfield_aux[lev][idim]->nGrowVect();
-                    auto pmf = std::make_unique<MultiFab>(Bfield_aux[lev][idim]->boxArray(),
-                                                                      dm, Bfield_aux[lev][idim]->nComp(), ng);
-                    // pmf->Redistribute(*Bfield_aux[lev][idim], 0, 0, Bfield_aux[lev][idim]->nComp(), ng);
-                    Bfield_aux[lev][idim] = std::move(pmf);
-                }
-                {
-                    const IntVect& ng = Efield_aux[lev][idim]->nGrowVect();
-                    auto pmf = std::make_unique<MultiFab>(Efield_aux[lev][idim]->boxArray(),
-                                                                      dm, Efield_aux[lev][idim]->nComp(), ng);
-                    // pmf->Redistribute(*Efield_aux[lev][idim], 0, 0, Efield_aux[lev][idim]->nComp(), ng);
-                    Efield_aux[lev][idim] = std::move(pmf);
-                }
+                RemakeMultiFab(Bfield_aux[lev][idim], dm, false);
+                RemakeMultiFab(Efield_aux[lev][idim], dm, false);
             }
         }
 
@@ -282,43 +259,12 @@ WarpX::RemakeLevel (int lev, Real /*time*/, const BoxArray& ba, const Distributi
         if (lev > 0) {
             for (int idim=0; idim < 3; ++idim)
             {
-                {
-                    const IntVect& ng = Bfield_cp[lev][idim]->nGrowVect();
-                    auto pmf = std::make_unique<MultiFab>(Bfield_cp[lev][idim]->boxArray(),
-                                                                      dm, Bfield_cp[lev][idim]->nComp(), ng);
-                    pmf->Redistribute(*Bfield_cp[lev][idim], 0, 0, Bfield_cp[lev][idim]->nComp(), ng);
-                    Bfield_cp[lev][idim] = std::move(pmf);
-                }
-                {
-                    const IntVect& ng = Efield_cp[lev][idim]->nGrowVect();
-                    auto pmf = std::make_unique<MultiFab>(Efield_cp[lev][idim]->boxArray(),
-                                                                      dm, Efield_cp[lev][idim]->nComp(), ng);
-                    pmf->Redistribute(*Efield_cp[lev][idim], 0, 0, Efield_cp[lev][idim]->nComp(), ng);
-                    Efield_cp[lev][idim] = std::move(pmf);
-                }
-                {
-                    const IntVect& ng = current_cp[lev][idim]->nGrowVect();
-                    auto pmf = std::make_unique<MultiFab>(current_cp[lev][idim]->boxArray(),
-                                                                       dm, current_cp[lev][idim]->nComp(), ng);
-                    current_cp[lev][idim] = std::move(pmf);
-                }
-            }
-
-            if (F_cp[lev] != nullptr) {
-                const IntVect& ng = F_cp[lev]->nGrowVect();
-                auto pmf = std::make_unique<MultiFab>(F_cp[lev]->boxArray(),
-                                                                  dm, F_cp[lev]->nComp(), ng);
-                pmf->Redistribute(*F_cp[lev], 0, 0, F_cp[lev]->nComp(), ng);
-                F_cp[lev] = std::move(pmf);
-            }
-
-            if (rho_cp[lev] != nullptr) {
-                const int nc = rho_cp[lev]->nComp();
-                const IntVect& ng = rho_cp[lev]->nGrowVect();
-                auto pmf  = std::make_unique<MultiFab>(rho_cp[lev]->boxArray(),
-                                                                  dm, nc, ng);
-                rho_cp[lev] = std::move(pmf);
+                RemakeMultiFab(Bfield_cp[lev][idim], dm, true);
+                RemakeMultiFab(Efield_cp[lev][idim], dm, true);
+                RemakeMultiFab(current_cp[lev][idim], dm, false);
             }
+            RemakeMultiFab(F_cp[lev], dm, true);
+            RemakeMultiFab(rho_cp[lev], dm, false);
 
 #ifdef WARPX_USE_PSATD
             if (maxwell_solver_id == MaxwellSolverAlgo::PSATD) {
@@ -358,57 +304,15 @@ WarpX::RemakeLevel (int lev, Real /*time*/, const BoxArray& ba, const Distributi
         if (lev > 0 && (n_field_gather_buffer > 0 || n_current_deposition_buffer > 0)) {
             for (int idim=0; idim < 3; ++idim)
             {
-                if (Bfield_cax[lev][idim])
-                {
-                    const IntVect& ng = Bfield_cax[lev][idim]->nGrowVect();
-                    auto pmf = std::make_unique<MultiFab>(Bfield_cax[lev][idim]->boxArray(),
-                                                                      dm, Bfield_cax[lev][idim]->nComp(), ng);
-                    // pmf->ParallelCopy(*Bfield_cax[lev][idim], 0, 0, Bfield_cax[lev][idim]->nComp(), ng, ng);
-                    Bfield_cax[lev][idim] = std::move(pmf);
-                }
-                if (Efield_cax[lev][idim])
-                {
-                    const IntVect& ng = Efield_cax[lev][idim]->nGrowVect();
-                    auto pmf = std::make_unique<MultiFab>(Efield_cax[lev][idim]->boxArray(),
-                                                                      dm, Efield_cax[lev][idim]->nComp(), ng);
-                    // pmf->ParallelCopy(*Efield_cax[lev][idim], 0, 0, Efield_cax[lev][idim]->nComp(), ng, ng);
-                    Efield_cax[lev][idim] = std::move(pmf);
-                }
-                if (current_buf[lev][idim])
-                {
-                    const IntVect& ng = current_buf[lev][idim]->nGrowVect();
-                    auto pmf = std::make_unique<MultiFab>(current_buf[lev][idim]->boxArray(),
-                                                                      dm, current_buf[lev][idim]->nComp(), ng);
-                    // pmf->ParallelCopy(*current_buf[lev][idim], 0, 0, current_buf[lev][idim]->nComp(), ng, ng);
-                    current_buf[lev][idim] = std::move(pmf);
-                }
-            }
-            if (charge_buf[lev])
-            {
-                const IntVect& ng = charge_buf[lev]->nGrowVect();
-                auto pmf = std::make_unique<MultiFab>(charge_buf[lev]->boxArray(),
-                                                                  dm, charge_buf[lev]->nComp(), ng);
-                // pmf->ParallelCopy(*charge_buf[lev][idim], 0, 0, charge_buf[lev]->nComp(), ng, ng);
-                charge_buf[lev] = std::move(pmf);
-            }
-            if (current_buffer_masks[lev])
-            {
-                const IntVect& ng = current_buffer_masks[lev]->nGrowVect();
-                auto pmf = std::make_unique<iMultiFab>(current_buffer_masks[lev]->boxArray(),
-                                                                    dm, current_buffer_masks[lev]->nComp(), ng);
-                // we can avoid this since we immediately re-build the values via BuildBufferMasks()
-                // pmf->Redistribute(*current_buffer_masks[lev], 0, 0, current_buffer_masks[lev]->nComp(), ng);
-                current_buffer_masks[lev] = std::move(pmf);
-            }
-            if (gather_buffer_masks[lev])
-            {
-                const IntVect& ng = gather_buffer_masks[lev]->nGrowVect();
-                auto pmf = std::make_unique<iMultiFab>(gather_buffer_masks[lev]->boxArray(),
-                                                                    dm, gather_buffer_masks[lev]->nComp(), ng);
-                // we can avoid this since we immediately re-build the values via BuildBufferMasks()
-                // pmf->Redistribute(*gather_buffer_masks[lev], 0, 0, gather_buffer_masks[lev]->nComp(), ng);
-                gather_buffer_masks[lev] = std::move(pmf);
+                RemakeMultiFab(Bfield_cax[lev][idim], dm, false);
+                RemakeMultiFab(Efield_cax[lev][idim], dm, false);
+                RemakeMultiFab(current_buf[lev][idim], dm, false);
             }
+            RemakeMultiFab(charge_buf[lev], dm, false);
+            // we can avoid redistributing these since we immediately re-build the values via BuildBufferMasks()
+            RemakeMultiFab(current_buffer_masks[lev], dm, false);
+            RemakeMultiFab(gather_buffer_masks[lev], dm, false);
+
             if (current_buffer_masks[lev] || gather_buffer_masks[lev])
                 BuildBufferMasks();
         }
@@ -458,7 +362,7 @@ WarpX::ComputeCostsHeuristic (amrex::Vector<std::unique_ptr<amrex::LayoutData<am
             }
         }
 
-        //Cell loop
+        // Cell loop
         MultiFab* Ex = Efield_fp[lev][0].get();
         for (MFIter mfi(*Ex, false); mfi.isValid(); ++mfi)
         {