From 368aa92d9f84101497909e088f853aa2f403a793 Mon Sep 17 00:00:00 2001
From: Rene Gassmoeller <rene.gassmoeller@mailbox.org>
Date: Sat, 21 Sep 2024 12:21:28 +0200
Subject: [PATCH] Simplify base weight sum

---
 include/aspect/particle/world.h |  6 ------
 source/particle/world.cc        | 37 ++++++++++++++++-----------------
 2 files changed, 18 insertions(+), 25 deletions(-)

diff --git a/include/aspect/particle/world.h b/include/aspect/particle/world.h
index 4cf7d466e1d..cbb9476d1ca 100644
--- a/include/aspect/particle/world.h
+++ b/include/aspect/particle/world.h
@@ -337,12 +337,6 @@ namespace aspect
          */
         typename ParticleLoadBalancing::Kind particle_load_balancing;
 
-        /**
-         * The index of this particle world. This is used to distinguish
-         * between multiple particle worlds in the same simulation.
-         */
-        unsigned int world_id;
-
         /**
          * Lower limit for particle number per cell. This limit is
          * useful for adaptive meshes to prevent fine cells from being empty
diff --git a/source/particle/world.cc b/source/particle/world.cc
index 3c98c691c2e..655b203a4e6 100644
--- a/source/particle/world.cc
+++ b/source/particle/world.cc
@@ -52,20 +52,6 @@ namespace aspect
     World<dim>::initialize()
     {
       CitationInfo::add("particles");
-      if (particle_load_balancing & ParticleLoadBalancing::repartition)
-        this->get_triangulation().signals.weight.connect(
-#if DEAL_II_VERSION_GTE(9,6,0)
-          [&] (const typename parallel::distributed::Triangulation<dim>::cell_iterator &cell,
-               const CellStatus status)
-          -> unsigned int
-#else
-          [&] (const typename parallel::distributed::Triangulation<dim>::cell_iterator &cell,
-               const typename parallel::distributed::Triangulation<dim>::CellStatus status)
-          -> unsigned int
-#endif
-        {
-          return this->cell_weight(cell, status);
-        });
 
       // Create a particle handler that stores the future particles.
       // If we restarted from a checkpoint we will fill this particle handler
@@ -373,9 +359,6 @@ namespace aspect
       if (cell->is_active() && !cell->is_locally_owned())
         return 0;
 
-      // Only add the base weight of cells in particle world 0, because the weight will
-      // be summed over all particle worlds.
-      const unsigned int base_weight = (world_id == 0) ? 1000 : 0;
       unsigned int n_particles_in_cell = 0;
       switch (status)
         {
@@ -410,7 +393,7 @@ namespace aspect
             Assert(false, ExcInternalError());
             break;
         }
-      return base_weight + n_particles_in_cell * particle_weight;
+      return n_particles_in_cell * particle_weight;
     }
 
 
@@ -834,7 +817,6 @@ namespace aspect
       // is not possible.
       const double CFL_number = prm.get_double ("CFL number");
       const unsigned int n_processes = Utilities::MPI::n_mpi_processes(this->get_mpi_communicator());
-      world_id = world_index;
 
       AssertThrow((n_processes == 1) || (CFL_number <= 1.0),
                   ExcMessage("The current particle algorithm does not work in "
@@ -899,6 +881,23 @@ namespace aspect
                                      "are listed in the corresponding manual subsection."));
           }
 
+        if (particle_load_balancing & ParticleLoadBalancing::repartition)
+          this->get_triangulation().signals.weight.connect(
+#if DEAL_II_VERSION_GTE(9,6,0)
+            [ &, world_index] (const typename parallel::distributed::Triangulation<dim>::cell_iterator &cell,
+                               const CellStatus status)
+            -> unsigned int
+#else
+            [ &, world_index] (const typename parallel::distributed::Triangulation<dim>::cell_iterator &cell,
+                               const typename parallel::distributed::Triangulation<dim>::CellStatus status)
+            -> unsigned int
+#endif
+          {
+            // Only add the base weight of cells in particle world 0, because all weights will be summed
+            // across all particle worlds.
+            return (world_index == 0) ? 1000 + this->cell_weight(cell, status) : this->cell_weight(cell, status);
+          });
+
 
         TimerOutput::Scope timer_section(this->get_computing_timer(), "Particles: Initialization");