From 350b959531a19749cb6d9c82f164b73599c08fed Mon Sep 17 00:00:00 2001
From: Hannah Klion <hannah.klion@gmail.com>
Date: Thu, 29 Feb 2024 14:52:34 -0800
Subject: [PATCH] use MakeSimilarDM and also attempt to average dowm the costs

---
 Source/Parallelization/WarpXRegrid.cpp | 38 +++++++++++++++++++-------
 Source/WarpX.H                         |  2 ++
 Source/WarpX.cpp                       |  1 +
 3 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/Source/Parallelization/WarpXRegrid.cpp b/Source/Parallelization/WarpXRegrid.cpp
index dcea3cab58a..17dfd6b93d6 100644
--- a/Source/Parallelization/WarpXRegrid.cpp
+++ b/Source/Parallelization/WarpXRegrid.cpp
@@ -19,6 +19,8 @@
 #include "Utils/WarpXAlgorithmSelection.H"
 #include "Utils/WarpXProfilerWrapper.H"
 
+#include <ablastr/coarsen/sample.H>
+
 #include <AMReX.H>
 #include <AMReX_BLassert.H>
 #include <AMReX_Box.H>
@@ -69,6 +71,12 @@ WarpX::LoadBalance ()
     int loadBalancedAnyLevel = false;
 
     const int nLevels = finestLevel();
+    if (do_similar_dm_refpatch) {
+        for (int lev = nLevels; lev > 0; --lev) {
+            ablastr::coarsen::sample::Coarsen(*costs[lev-1], *costs[lev],0,0,1,0,WarpX::RefRatio(lev-1));
+        }
+    }
+
     for (int lev = 0; lev <= nLevels; ++lev)
     {
         int doLoadBalance = false;
@@ -83,16 +91,26 @@ WarpX::LoadBalance ()
         amrex::Real currentEfficiency = 0.0;
         amrex::Real proposedEfficiency = 0.0;
 
-        newdm = (load_balance_with_sfc)
-            ? DistributionMapping::makeSFC(*costs[lev],
-                                           currentEfficiency, proposedEfficiency,
-                                           false,
-                                           ParallelDescriptor::IOProcessorNumber())
-            : DistributionMapping::makeKnapSack(*costs[lev],
-                                                currentEfficiency, proposedEfficiency,
-                                                nmax,
-                                                false,
-                                                ParallelDescriptor::IOProcessorNumber());
+        if (lev == 0 || !do_similar_dm_refpatch) {
+            newdm = (load_balance_with_sfc)
+                ? DistributionMapping::makeSFC(*costs[lev],
+                                               currentEfficiency, proposedEfficiency,
+                                               false,
+                                               ParallelDescriptor::IOProcessorNumber())
+                : DistributionMapping::makeKnapSack(*costs[lev],
+                                                    currentEfficiency, proposedEfficiency,
+                                                    nmax,
+                                                    false,
+                                                    ParallelDescriptor::IOProcessorNumber());
+        } else {
+            amrex::BoxArray coarse_ba = boxArray(lev-1);
+            amrex::DistributionMapping coarse_dm = DistributionMap(lev-1);
+            amrex::BoxArray ba = boxArray(lev);
+            ba.coarsen(WarpX::RefRatio(lev-1));
+            newdm = amrex::MakeSimilarDM(ba, coarse_ba, coarse_dm, getngEB());
+        }
+        Print() << "new dm on lev " << lev << ": \n";
+        Print() << newdm << std::endl;
         // As specified in the above calls to makeSFC and makeKnapSack, the new
         // distribution mapping is NOT communicated to all ranks; the loadbalanced
         // dm is up-to-date only on root, and we can decide whether to broadcast
diff --git a/Source/WarpX.H b/Source/WarpX.H
index 6526171ff67..d24ab99e6aa 100644
--- a/Source/WarpX.H
+++ b/Source/WarpX.H
@@ -1641,6 +1641,8 @@ private:
      * uniform plasma on a domain of size 128 by 128 by 128, from which the approximate
      * time per iteration per particle is computed. */
     amrex::Real costs_heuristic_particles_wt = amrex::Real(0);
+    /** Whether to use MakeSimilarDM when load balancing lev > 1 */
+    int do_similar_dm_refpatch = 0;
 
     // Determines timesteps for override sync
     utils::parser::IntervalsParser override_sync_intervals;
diff --git a/Source/WarpX.cpp b/Source/WarpX.cpp
index c929718b043..272e40ef0c4 100644
--- a/Source/WarpX.cpp
+++ b/Source/WarpX.cpp
@@ -1307,6 +1307,7 @@ WarpX::ReadParameters ()
         load_balance_intervals = utils::parser::IntervalsParser(
             load_balance_intervals_string_vec);
         pp_algo.query("load_balance_with_sfc", load_balance_with_sfc);
+        pp_algo.query("do_similar_dm_refpatch", do_similar_dm_refpatch);
         // Knapsack factor only used with non-SFC strategy
         if (!load_balance_with_sfc) {
             pp_algo.query("load_balance_knapsack_factor", load_balance_knapsack_factor);