2 files changed, 29 insertions, 0 deletions
diff --git a/Source/Utils/WarpXUtil.H b/Source/Utils/WarpXUtil.H
index 61952da17..dfbed3e82 100644
--- a/Source/Utils/WarpXUtil.H
+++ b/Source/Utils/WarpXUtil.H
@@ -8,8 +8,11 @@
 #ifndef WARPX_UTILS_H_
 #define WARPX_UTILS_H_
 
+#include <AMReX_BoxArray.H>
+#include <AMReX_DistributionMapping.H>
 #include <AMReX_Extension.H>
 #include <AMReX_GpuQualifiers.H>
+#include <AMReX_LayoutData.H>
 #include <AMReX_ParmParse.H>
 #include <AMReX_Parser.H>
 #include <AMReX_REAL.H>
@@ -397,4 +400,18 @@ namespace WarpXUtilStr
 
 }
 
+namespace WarpXUtilLoadBalance
+{
+    /** \brief We only want to update the cost data if the grids we are working on
+     *  are the main grids, i.e. not the PML grids. This function returns whether
+     *   this is the case or not.
+     * @param[in] cost pointer to the cost data
+     * @param[in] ba the grids to check
+     * @param[in] dm the dmap to check
+     * @return consistent whether the grids are consistent or not.
+     */
+    bool doCosts (const amrex::LayoutData<amrex::Real>* cost, const amrex::BoxArray ba,
+                  const amrex::DistributionMapping& dm);
+}
+
 #endif //WARPX_UTILS_H_
diff --git a/Source/Utils/WarpXUtil.cpp b/Source/Utils/WarpXUtil.cpp
index 03b3768a7..f7b89ee0d 100644
--- a/Source/Utils/WarpXUtil.cpp
+++ b/Source/Utils/WarpXUtil.cpp
@@ -740,3 +740,15 @@ namespace WarpXUtilStr
     }
 
 }
+
+namespace WarpXUtilLoadBalance
+{
+    bool doCosts (const amrex::LayoutData<amrex::Real>* costs, const amrex::BoxArray ba,
+                  const amrex::DistributionMapping& dm)
+    {
+        bool consistent = costs && (dm == costs->DistributionMap()) &&
+            (ba.CellEqual(costs->boxArray())) &&
+            (WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers);
+        return consistent;
+    }
+}