cleaning

lucafedeli88 · Jul 29, 2024 · 9a86c6f · 9a86c6f
1 parent a144b99
commit 9a86c6f
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 7 deletions.
diff --git a/Source/Diagnostics/ReducedDiags/MultiReducedDiags.H b/Source/Diagnostics/ReducedDiags/MultiReducedDiags.H
@@ -53,6 +53,10 @@ public:
      *  @param[in] step current iteration time */
     void WriteToFile (int step);
 
+    bool MustRecomputeCosts(int step);
+
+private:
+    std::vector<size_t> m_load_balance_diags;
 };
 
 #endif
diff --git a/Source/Diagnostics/ReducedDiags/MultiReducedDiags.cpp b/Source/Diagnostics/ReducedDiags/MultiReducedDiags.cpp
@@ -71,7 +71,7 @@ MultiReducedDiags::MultiReducedDiags ()
     };
     // loop over all reduced diags and fill m_multi_rd with requested reduced diags
     std::transform(m_rd_names.begin(), m_rd_names.end(), std::back_inserter(m_multi_rd),
-        [&](const auto& rd_name){
+        [&, idx = 0](const auto& rd_name) mutable {
             const ParmParse pp_rd_name(rd_name);
 
             // read reduced diags type
@@ -83,6 +83,10 @@ MultiReducedDiags::MultiReducedDiags ()
                 rd_type + " is not a valid type for reduced diagnostic " + rd_name
             );
 
+            if (rd_type == "LoadBalanceCosts" || "LoadBalanceEfficiency"){
+                m_load_balance_diags.push_back(idx++);
+            }
+
             return reduced_diags_dictionary.at(rd_type)(rd_name);
         });
     // end loop over all reduced diags
@@ -137,4 +141,15 @@ void MultiReducedDiags::WriteToFile (int step)
     }
     // end loop over all reduced diags
 }
+
+bool MultiReducedDiags::MustRecomputeCosts(int step)
+{
+    for (const auto idx : m_load_balance_diags){
+        if ( m_multi_rd[idx]->m_intervals.contains(step+1)){
+            return true;
+        }
+    }
+    return false;
+}
+
 // end void MultiReducedDiags::WriteToFile
diff --git a/Source/Evolve/WarpXEvolve.cpp b/Source/Evolve/WarpXEvolve.cpp
@@ -22,6 +22,7 @@
 #       include "FieldSolver/SpectralSolver/SpectralSolver.H"
 #   endif
 #endif
+#include "LoadBalance/LoadBalance.H"
 #include "Parallelization/GuardCellManager.H"
 #include "Particles/MultiParticleContainer.H"
 #include "Fluids/MultiFluidContainer.H"
@@ -58,6 +59,7 @@
 
 using namespace amrex;
 using ablastr::utils::SignalHandling;
+using namespace warpx::load_balance;
 
 void
 WarpX::Evolve (int numsteps)
@@ -262,11 +264,13 @@ WarpX::Evolve (int numsteps)
         ExecutePythonCallback("afterstep");
 
         /// reduced diags
-        //
-        //  TODO: conditionally recompute weights
-        //
         if (reduced_diags->m_plot_rd != 0)
         {
+            if (reduced_diags->MustRecomputeCosts(step)){
+                auto& load_balance = LoadBalance::get_instance();
+                load_balance.compute_costs_if_heuristic(
+                    finestLevel(), Efield_fp, *mypc);
+            }
             reduced_diags->LoadBalance();
             reduced_diags->ComputeDiags(step);
             reduced_diags->WriteToFile(step);

diff --git a/Source/Initialization/WarpXInitData.cpp b/Source/Initialization/WarpXInitData.cpp
@@ -531,11 +531,13 @@ WarpX::InitData ()
         multi_diags->FilterComputePackFlush(istep[0] - 1);
 
         // Write reduced diagnostics before the first iteration.
-        //
-        //  TODO: conditionally recompute weights
-        //
         if (reduced_diags->m_plot_rd != 0)
         {
+            if (reduced_diags->MustRecomputeCosts(istep[0]-1)){
+                auto& load_balance = LoadBalance::get_instance();
+                load_balance.compute_costs_if_heuristic(
+                    finestLevel(), Efield_fp, *mypc);
+            }
             reduced_diags->ComputeDiags(istep[0] - 1);
             reduced_diags->WriteToFile(istep[0] - 1);
         }