From 3c7cb6a4467bb6ce73c045cb3e63e12e3d40178e Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Thu, 9 Oct 2025 22:11:45 +0200
Subject: [PATCH 01/35] weight extension with dummy weights

---
 src/t8_forest/t8_forest_partition.cxx | 56 ++++++++++++++++++++++++---
 1 file changed, 51 insertions(+), 5 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index fd396d77fd..e490b371e4 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -32,7 +32,7 @@ T8_EXTERN_C_BEGIN ();
 
 /**
  * For each tree that we send elements from to other processes,
- * we send the information stored in this struct to the other process 
+ * we send the information stored in this struct to the other process
  */
 typedef struct
 {
@@ -202,7 +202,7 @@ t8_forest_partition_test_boundary_element ([[maybe_unused]] const t8_forest_t fo
     return;
   }
   if (forest->mpirank == forest->mpisize - 1) {
-    /* The last process can only share a tree with process rank-1. 
+    /* The last process can only share a tree with process rank-1.
      * However, this is already tested by process rank-1. */
     return;
   }
@@ -211,7 +211,7 @@ t8_forest_partition_test_boundary_element ([[maybe_unused]] const t8_forest_t fo
     /* The first tree on process rank+1 is not shared with current rank, nothing to do */
     return;
   }
-  /* The first tree on process rank+1 may be shared but empty. 
+  /* The first tree on process rank+1 may be shared but empty.
    * Thus, the first descendant id of rank+1 is not of the first local tree. */
   if (local_tree_num_elements == 0) {
     /* check if first not shared tree of process rank+1 contains elements */
@@ -428,12 +428,58 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   mpiret = sc_MPI_Comm_size (comm, &mpisize);
   SC_CHECK_MPI (mpiret);
 
+  auto weight_fcn = []( auto&&... ) -> double { return 1.; }; // should give same result as before
+
+  double const partition_weight = [&](){
+    double retval = 0.;
+    for (int ltreeid = 0; ltreeid < t8_forest_get_num_local_trees(forest); ++ltreeid){
+      for ( int elm_in_tree = 0; elm_in_tree < t8_forest_get_tree_num_leaf_elements(forest, ltreeid); ++elm_in_tree ){
+        retval += weight_fcn( forest, ltreeid, elm_in_tree );
+      }
+    }
+    return retval;
+  }();
+
+  double const partition_weight_offset = [&](){
+    double retval = 0.;
+    sc_MPI_Scan(&partition_weight, &retval, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
+    return retval;
+  }();
+
+  double const target_weight = [&](){
+    double total_forest_weight = partition_weight_offset + partition_weight;
+    sc_MPI_Bcast(&total_forest_weight, 1, sc_MPI_DOUBLE, mpisize-1, forest->mpicomm);
+    return total_forest_weight / mpisize;
+  }();
+
+  t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id(forest);
+
   if (t8_shmem_array_start_writing (forest->element_offsets)) {
     if (forest_from->global_num_leaf_elements > 0) {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
+
+      t8_locidx_t current_tree = 0;
+      t8_locidx_t current_elm_in_tree = 0;
+      double accumulated_weight = partition_weight_offset;
       for (i = 0; i < mpisize; i++) {
-        /* Calculate the first element index for each process. We convert to doubles to prevent overflow */
-        new_first_element_id = (((double) i * (long double) forest_from->global_num_leaf_elements) / (double) mpisize);
+        if (i*target_weight < partition_weight_offset or
+            i*target_weight > partition_weight_offset + partition_weight) {
+          continue; // the new first element for the i-th partition is not here
+        }
+        while (accumulated_weight < i*target_weight){
+          T8_ASSERT(current_tree < t8_forest_get_num_local_trees(forest));
+          T8_ASSERT(current_elm_in_tree < t8_forest_get_tree_num_leaf_elements(forest, current_tree));
+          accumulated_weight += weight_fcn( forest, current_tree, current_elm_in_tree );
+          ++current_elm_in_tree;
+          if ( current_elm_in_tree == t8_forest_get_tree_num_leaf_elements(forest, current_tree)){
+            ++current_tree;
+            current_elm_in_tree = 0;
+          }
+        }
+        /* Calculate the first element index for the i-th process */
+        new_first_element_id = partition_offset
+                             + t8_forest_get_tree_element_offset(forest, current_tree)
+                             + current_elm_in_tree;
         T8_ASSERT (0 <= new_first_element_id && new_first_element_id < forest_from->global_num_leaf_elements);
         element_offsets[i] = new_first_element_id;
       }

From ebe78e5f282e6f1f973f8a59fdb1168dd1756ca7 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Thu, 9 Oct 2025 22:44:34 +0200
Subject: [PATCH 02/35] cleanup and add comments

---
 src/t8_forest/t8_forest_partition.cxx | 40 ++++++++++++++-------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index e490b371e4..8193c2b413 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -408,16 +408,12 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
 static void
 t8_forest_partition_compute_new_offset (t8_forest_t forest)
 {
-  t8_forest_t forest_from;
-  sc_MPI_Comm comm;
-  t8_gloidx_t new_first_element_id;
-  int i, mpiret, mpisize;
-
   T8_ASSERT (t8_forest_is_initialized (forest));
   T8_ASSERT (forest->set_from != NULL);
 
-  forest_from = forest->set_from;
-  comm = forest->mpicomm;
+  t8_forest_t forest_from = forest->set_from;
+  sc_MPI_Comm comm = forest->mpicomm;
+  int mpisize = forest->mpisize;
 
   T8_ASSERT (forest->element_offsets == NULL);
   /* Set the shmem array type to comm */
@@ -425,33 +421,37 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   t8_shmem_set_type (comm, T8_SHMEM_BEST_TYPE);
   /* Initialize the shmem array */
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
-  mpiret = sc_MPI_Comm_size (comm, &mpisize);
-  SC_CHECK_MPI (mpiret);
 
   auto weight_fcn = []( auto&&... ) -> double { return 1.; }; // should give same result as before
 
+  // Compute the weight of all the partition-local elements as a whole
   double const partition_weight = [&](){
     double retval = 0.;
-    for (int ltreeid = 0; ltreeid < t8_forest_get_num_local_trees(forest); ++ltreeid){
-      for ( int elm_in_tree = 0; elm_in_tree < t8_forest_get_tree_num_leaf_elements(forest, ltreeid); ++elm_in_tree ){
-        retval += weight_fcn( forest, ltreeid, elm_in_tree );
+    // t8_forest_get_num_local_trees & t8_forest_get_tree_num_leaf_elements assume a
+    // commited forest, which is not the case here yet...
+    for (int ltreeid = 0; ltreeid < t8_forest_get_num_local_trees(forest); ++ltreeid) {
+      for ( int ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements(forest, ltreeid); ++ielm ){
+        retval += weight_fcn( forest, ltreeid, ielm ); // first pass
       }
     }
     return retval;
   }();
 
+  // Compute the aggregated weight of all the partitions of lower rank
   double const partition_weight_offset = [&](){
     double retval = 0.;
     sc_MPI_Scan(&partition_weight, &retval, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
     return retval;
   }();
 
+  // Compute the total weight of the forest and derive the equilibrium weight for each partition
   double const target_weight = [&](){
     double total_forest_weight = partition_weight_offset + partition_weight;
     sc_MPI_Bcast(&total_forest_weight, 1, sc_MPI_DOUBLE, mpisize-1, forest->mpicomm);
     return total_forest_weight / mpisize;
   }();
 
+  // t8_forest_get_first_local_leaf_element_id too requires a commited forest I believe...
   t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id(forest);
 
   if (t8_shmem_array_start_writing (forest->element_offsets)) {
@@ -461,15 +461,17 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
       t8_locidx_t current_tree = 0;
       t8_locidx_t current_elm_in_tree = 0;
       double accumulated_weight = partition_weight_offset;
-      for (i = 0; i < mpisize; i++) {
+      for (int i = 0; i < mpisize; i++) {
         if (i*target_weight < partition_weight_offset or
             i*target_weight > partition_weight_offset + partition_weight) {
           continue; // the new first element for the i-th partition is not here
         }
-        while (accumulated_weight < i*target_weight){
+        // this while loop is a very ugly way to say "find the element where the partial sum
+        // of all the weights up to this point matches i*target_weight"
+        while (accumulated_weight < i*target_weight) {
           T8_ASSERT(current_tree < t8_forest_get_num_local_trees(forest));
           T8_ASSERT(current_elm_in_tree < t8_forest_get_tree_num_leaf_elements(forest, current_tree));
-          accumulated_weight += weight_fcn( forest, current_tree, current_elm_in_tree );
+          accumulated_weight += weight_fcn( forest, current_tree, current_elm_in_tree ); // second pass
           ++current_elm_in_tree;
           if ( current_elm_in_tree == t8_forest_get_tree_num_leaf_elements(forest, current_tree)){
             ++current_tree;
@@ -477,9 +479,9 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
           }
         }
         /* Calculate the first element index for the i-th process */
-        new_first_element_id = partition_offset
-                             + t8_forest_get_tree_element_offset(forest, current_tree)
-                             + current_elm_in_tree;
+        t8_gloidx_t new_first_element_id = partition_offset
+                                         + t8_forest_get_tree_element_offset(forest, current_tree) // same problem here
+                                         + current_elm_in_tree;
         T8_ASSERT (0 <= new_first_element_id && new_first_element_id < forest_from->global_num_leaf_elements);
         element_offsets[i] = new_first_element_id;
       }
@@ -487,7 +489,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
     }
     else {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
-      for (i = 0; i <= mpisize; i++) {
+      for (int i = 0; i <= mpisize; i++) {
         element_offsets[i] = 0;
       }
     }

From 8c4170b01591bb2ce285826da44f146a4db53fac Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Fri, 10 Oct 2025 11:01:55 +0200
Subject: [PATCH 03/35] fix scan bug

---
 src/t8_forest/t8_forest_partition.cxx | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 8193c2b413..9ed86fca12 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -441,6 +441,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   double const partition_weight_offset = [&](){
     double retval = 0.;
     sc_MPI_Scan(&partition_weight, &retval, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
+    retval -= partition_weight; // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
     return retval;
   }();
 

From 27a54b3c7e34515d29c1d62002b5f8c7a729c40e Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Fri, 10 Oct 2025 11:03:52 +0200
Subject: [PATCH 04/35] indent

---
 src/t8_forest/t8_forest_partition.cxx | 44 +++++++++++++--------------
 1 file changed, 22 insertions(+), 22 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 9ed86fca12..78414dd75d 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -422,38 +422,38 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   /* Initialize the shmem array */
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
 
-  auto weight_fcn = []( auto&&... ) -> double { return 1.; }; // should give same result as before
+  auto weight_fcn = [] (auto &&...) -> double { return 1.; };  // should give same result as before
 
   // Compute the weight of all the partition-local elements as a whole
-  double const partition_weight = [&](){
+  double const partition_weight = [&] () {
     double retval = 0.;
     // t8_forest_get_num_local_trees & t8_forest_get_tree_num_leaf_elements assume a
     // commited forest, which is not the case here yet...
-    for (int ltreeid = 0; ltreeid < t8_forest_get_num_local_trees(forest); ++ltreeid) {
-      for ( int ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements(forest, ltreeid); ++ielm ){
-        retval += weight_fcn( forest, ltreeid, ielm ); // first pass
+    for (int ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
+      for (int ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
+        retval += weight_fcn (forest, ltreeid, ielm);  // first pass
       }
     }
     return retval;
   }();
 
   // Compute the aggregated weight of all the partitions of lower rank
-  double const partition_weight_offset = [&](){
+  double const partition_weight_offset = [&] () {
     double retval = 0.;
-    sc_MPI_Scan(&partition_weight, &retval, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
-    retval -= partition_weight; // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
+    sc_MPI_Scan (&partition_weight, &retval, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
+    retval -= partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
     return retval;
   }();
 
   // Compute the total weight of the forest and derive the equilibrium weight for each partition
-  double const target_weight = [&](){
+  double const target_weight = [&] () {
     double total_forest_weight = partition_weight_offset + partition_weight;
-    sc_MPI_Bcast(&total_forest_weight, 1, sc_MPI_DOUBLE, mpisize-1, forest->mpicomm);
+    sc_MPI_Bcast (&total_forest_weight, 1, sc_MPI_DOUBLE, mpisize - 1, forest->mpicomm);
     return total_forest_weight / mpisize;
   }();
 
   // t8_forest_get_first_local_leaf_element_id too requires a commited forest I believe...
-  t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id(forest);
+  t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id (forest);
 
   if (t8_shmem_array_start_writing (forest->element_offsets)) {
     if (forest_from->global_num_leaf_elements > 0) {
@@ -463,26 +463,26 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
       t8_locidx_t current_elm_in_tree = 0;
       double accumulated_weight = partition_weight_offset;
       for (int i = 0; i < mpisize; i++) {
-        if (i*target_weight < partition_weight_offset or
-            i*target_weight > partition_weight_offset + partition_weight) {
-          continue; // the new first element for the i-th partition is not here
+        if (i * target_weight < partition_weight_offset
+            or i * target_weight > partition_weight_offset + partition_weight) {
+          continue;  // the new first element for the i-th partition is not here
         }
         // this while loop is a very ugly way to say "find the element where the partial sum
         // of all the weights up to this point matches i*target_weight"
-        while (accumulated_weight < i*target_weight) {
-          T8_ASSERT(current_tree < t8_forest_get_num_local_trees(forest));
-          T8_ASSERT(current_elm_in_tree < t8_forest_get_tree_num_leaf_elements(forest, current_tree));
-          accumulated_weight += weight_fcn( forest, current_tree, current_elm_in_tree ); // second pass
+        while (accumulated_weight < i * target_weight) {
+          T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest));
+          T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest, current_tree));
+          accumulated_weight += weight_fcn (forest, current_tree, current_elm_in_tree);  // second pass
           ++current_elm_in_tree;
-          if ( current_elm_in_tree == t8_forest_get_tree_num_leaf_elements(forest, current_tree)){
+          if (current_elm_in_tree == t8_forest_get_tree_num_leaf_elements (forest, current_tree)) {
             ++current_tree;
             current_elm_in_tree = 0;
           }
         }
         /* Calculate the first element index for the i-th process */
-        t8_gloidx_t new_first_element_id = partition_offset
-                                         + t8_forest_get_tree_element_offset(forest, current_tree) // same problem here
-                                         + current_elm_in_tree;
+        t8_gloidx_t new_first_element_id
+          = partition_offset + t8_forest_get_tree_element_offset (forest, current_tree)  // same problem here
+            + current_elm_in_tree;
         T8_ASSERT (0 <= new_first_element_id && new_first_element_id < forest_from->global_num_leaf_elements);
         element_offsets[i] = new_first_element_id;
       }

From 3d5c2b6f5e2a67946a329ff68e39185a5d00e7ba Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Fri, 10 Oct 2025 11:34:19 +0200
Subject: [PATCH 05/35] expose the weight function as an optional parameter

---
 src/t8_forest/t8_forest_partition.cxx | 18 ++++++++++--------
 src/t8_forest/t8_forest_partition.h   | 19 ++++++++++++++-----
 2 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 78414dd75d..fd76db080b 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -406,7 +406,7 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
 /* Calculate the new element_offset for forest from
  * the element in forest->set_from assuming a partition without element weights */
 static void
-t8_forest_partition_compute_new_offset (t8_forest_t forest)
+t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t* weight_fcn)
 {
   T8_ASSERT (t8_forest_is_initialized (forest));
   T8_ASSERT (forest->set_from != NULL);
@@ -422,16 +422,18 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   /* Initialize the shmem array */
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
 
-  auto weight_fcn = [] (auto &&...) -> double { return 1.; };  // should give same result as before
+  if (weight_fcn == nullptr){
+    weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1.; };
+  }
 
   // Compute the weight of all the partition-local elements as a whole
   double const partition_weight = [&] () {
     double retval = 0.;
     // t8_forest_get_num_local_trees & t8_forest_get_tree_num_leaf_elements assume a
     // commited forest, which is not the case here yet...
-    for (int ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
-      for (int ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
-        retval += weight_fcn (forest, ltreeid, ielm);  // first pass
+    for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
+      for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
+        retval += weight_fcn (forest, ltreeid, ielm);
       }
     }
     return retval;
@@ -472,7 +474,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
         while (accumulated_weight < i * target_weight) {
           T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest));
           T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest, current_tree));
-          accumulated_weight += weight_fcn (forest, current_tree, current_elm_in_tree);  // second pass
+          accumulated_weight += weight_fcn (forest, current_tree, current_elm_in_tree);
           ++current_elm_in_tree;
           if (current_elm_in_tree == t8_forest_get_tree_num_leaf_elements (forest, current_tree)) {
             ++current_tree;
@@ -1218,7 +1220,7 @@ t8_forest_partition_given (t8_forest_t forest, const int send_data, const sc_arr
  * Currently the elements are distributed evenly (each element has the same weight).
  */
 void
-t8_forest_partition (t8_forest_t forest)
+t8_forest_partition (t8_forest_t forest, weight_fcn_t* weight_callback)
 {
   t8_forest_t forest_from;
   int create_offset_from = 0;
@@ -1248,7 +1250,7 @@ t8_forest_partition (t8_forest_t forest)
   /* TODO: if offsets already exist on forest_from, check it for consistency */
 
   /* We now calculate the new element offsets */
-  t8_forest_partition_compute_new_offset (forest);
+  t8_forest_partition_compute_new_offset (forest, weight_callback);
   t8_forest_partition_given (forest, 0, NULL, NULL);
 
   T8_ASSERT ((size_t) t8_forest_get_num_local_trees (forest_from) == forest_from->trees->elem_count);
diff --git a/src/t8_forest/t8_forest_partition.h b/src/t8_forest/t8_forest_partition.h
index d5da1ab8b1..00f1d33d5c 100644
--- a/src/t8_forest/t8_forest_partition.h
+++ b/src/t8_forest/t8_forest_partition.h
@@ -32,15 +32,24 @@
 #include <t8.h>
 #include <t8_forest/t8_forest_general.h>
 
+/**
+ * The prototype a weight function for the t8_forest_partition algorithm.
+ * The function should be pure, and return a positive weight given a forest, a local tree index and an element index within the local tree
+ */
+using weight_fcn_t = double(t8_forest_t, t8_locidx_t, t8_locidx_t);
+
 T8_EXTERN_C_BEGIN ();
 /**
  * Populate a forest with the partitioned elements of forest->set_from.
  * Currently the elements are distributed evenly (each element has the same weight).
- * 
+ *
  * \param [in,out]  forest  The forest.
+ * \param [in]      weight_callback A callback function defining element weights for the partitioning
+ * \pre \a weight_callback must be free of side-effects, the behavior is undefined otherwise
+ * \note If \a weight_callback is null, then all the elements are assumed to have the same weight
 */
 void
-t8_forest_partition (t8_forest_t forest);
+t8_forest_partition (t8_forest_t forest, weight_fcn_t* weight_callback = nullptr);
 
 /** Create the element_offset array of a partitioned forest.
  * \param [in,out]  forest The forest.
@@ -76,13 +85,13 @@ t8_forest_partition_create_first_desc (t8_forest_t forest);
 void
 t8_forest_partition_create_tree_offsets (t8_forest_t forest);
 
-/** \brief Re-Partition an array accordingly to a partitioned forest. 
- * 
+/** \brief Re-Partition an array accordingly to a partitioned forest.
+ *
  * \param[in] forest_from The forest before the partitioning step.
  * \param[in] forest_to The partitioned forest of \a forest_from.
  * \param[in] data_in A pointer to an sc_array_t holding data (one value per element) accordingly to \a forest_from.
  * \param[in,out] data_out A pointer to an already allocated sc_array_t capable of holding data accordingly to \a forest_to.
- * 
+ *
  * \note \a data_in has to be of size equal to the number of local elements of \a forest_from
  * \a data_out has to be already allocated and has to be of size equal to the number of local elements of \a forest_to.
 */

From 0d6ab35b167837cfd1fa04b352a0c8f7b7c6ee44 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Fri, 10 Oct 2025 11:36:14 +0200
Subject: [PATCH 06/35] update comment and indent

---
 src/t8_forest/t8_forest_partition.cxx | 8 ++++----
 src/t8_forest/t8_forest_partition.h   | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index fd76db080b..efe70e0e67 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -404,9 +404,9 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
 }
 
 /* Calculate the new element_offset for forest from
- * the element in forest->set_from assuming a partition without element weights */
+ * the element in forest->set_from using the provided weight function */
 static void
-t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t* weight_fcn)
+t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight_fcn)
 {
   T8_ASSERT (t8_forest_is_initialized (forest));
   T8_ASSERT (forest->set_from != NULL);
@@ -422,7 +422,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t* weight
   /* Initialize the shmem array */
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
 
-  if (weight_fcn == nullptr){
+  if (weight_fcn == nullptr) {
     weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1.; };
   }
 
@@ -1220,7 +1220,7 @@ t8_forest_partition_given (t8_forest_t forest, const int send_data, const sc_arr
  * Currently the elements are distributed evenly (each element has the same weight).
  */
 void
-t8_forest_partition (t8_forest_t forest, weight_fcn_t* weight_callback)
+t8_forest_partition (t8_forest_t forest, weight_fcn_t *weight_callback)
 {
   t8_forest_t forest_from;
   int create_offset_from = 0;
diff --git a/src/t8_forest/t8_forest_partition.h b/src/t8_forest/t8_forest_partition.h
index 00f1d33d5c..0f96605239 100644
--- a/src/t8_forest/t8_forest_partition.h
+++ b/src/t8_forest/t8_forest_partition.h
@@ -36,7 +36,7 @@
  * The prototype a weight function for the t8_forest_partition algorithm.
  * The function should be pure, and return a positive weight given a forest, a local tree index and an element index within the local tree
  */
-using weight_fcn_t = double(t8_forest_t, t8_locidx_t, t8_locidx_t);
+using weight_fcn_t = double (t8_forest_t, t8_locidx_t, t8_locidx_t);
 
 T8_EXTERN_C_BEGIN ();
 /**
@@ -49,7 +49,7 @@ T8_EXTERN_C_BEGIN ();
  * \note If \a weight_callback is null, then all the elements are assumed to have the same weight
 */
 void
-t8_forest_partition (t8_forest_t forest, weight_fcn_t* weight_callback = nullptr);
+t8_forest_partition (t8_forest_t forest, weight_fcn_t *weight_callback = nullptr);
 
 /** Create the element_offset array of a partitioned forest.
  * \param [in,out]  forest The forest.

From 9075a91bb12378cd5834e802b66bdfdc3ccc0fef Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 14 Oct 2025 13:47:51 +0200
Subject: [PATCH 07/35] add nullptr fast-forwards

---
 src/t8_forest/t8_forest_partition.cxx | 34 ++++++++++++++++++---------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index efe70e0e67..306efa1b31 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -422,12 +422,11 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
   /* Initialize the shmem array */
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
 
-  if (weight_fcn == nullptr) {
-    weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1.; };
-  }
-
   // Compute the weight of all the partition-local elements as a whole
-  double const partition_weight = [&] () {
+  double const partition_weight = [&] () -> double {
+    if (weight_fcn == nullptr) {
+      return t8_forest_get_local_num_leaf_elements (forest);
+    }
     double retval = 0.;
     // t8_forest_get_num_local_trees & t8_forest_get_tree_num_leaf_elements assume a
     // commited forest, which is not the case here yet...
@@ -440,7 +439,10 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
   }();
 
   // Compute the aggregated weight of all the partitions of lower rank
-  double const partition_weight_offset = [&] () {
+  double const partition_weight_offset = [&] () -> double {
+    if (weight_fcn == nullptr) {
+      return t8_forest_get_first_local_leaf_element_id (forest);
+    }
     double retval = 0.;
     sc_MPI_Scan (&partition_weight, &retval, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
     retval -= partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
@@ -448,12 +450,17 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
   }();
 
   // Compute the total weight of the forest and derive the equilibrium weight for each partition
-  double const target_weight = [&] () {
-    double total_forest_weight = partition_weight_offset + partition_weight;
-    sc_MPI_Bcast (&total_forest_weight, 1, sc_MPI_DOUBLE, mpisize - 1, forest->mpicomm);
-    return total_forest_weight / mpisize;
+  double const total_forest_weight = [&] () -> double {
+    if (weight_fcn == nullptr) {
+      return t8_forest_get_global_num_leaf_elements (forest);
+    }
+    double retval = partition_weight_offset + partition_weight;
+    sc_MPI_Bcast (&retval, 1, sc_MPI_DOUBLE, mpisize - 1, forest->mpicomm);
+    return retval;
   }();
 
+  double const target_weight = total_forest_weight / mpisize;
+
   // t8_forest_get_first_local_leaf_element_id too requires a commited forest I believe...
   t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id (forest);
 
@@ -474,7 +481,12 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
         while (accumulated_weight < i * target_weight) {
           T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest));
           T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest, current_tree));
-          accumulated_weight += weight_fcn (forest, current_tree, current_elm_in_tree);
+          if (weight_fcn == nullptr) {
+            accumulated_weight += 1;
+          }
+          else {
+            accumulated_weight += weight_fcn (forest, current_tree, current_elm_in_tree);
+          }
           ++current_elm_in_tree;
           if (current_elm_in_tree == t8_forest_get_tree_num_leaf_elements (forest, current_tree)) {
             ++current_tree;

From b17975f0a751a9cd53846f4aa8aabef7d6e88689 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 14 Oct 2025 15:12:29 +0200
Subject: [PATCH 08/35] extract first pass into helper function

---
 src/t8_forest/t8_forest_partition.cxx | 87 ++++++++++++++-------------
 1 file changed, 45 insertions(+), 42 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 306efa1b31..46e308429b 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -403,6 +403,45 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
   }
 }
 
+// Computes the weight of the forest, the weight of the local partition, and the local offset (i.e. the weight
+// of all the partitions of lower rank combined)
+// If weight_fcn is null, all the elements are assumed to be of unit weight.
+static std::tuple< double, double, double >
+t8_forest_integrate_leaf_weights (t8_forest_t forest, weight_fcn_t *weight_fcn)
+{
+    T8_ASSERT( t8_forest_is_committed(forest) );
+
+    if (weight_fcn == nullptr) {
+      return {
+        t8_forest_get_global_num_leaf_elements (forest),
+        t8_forest_get_local_num_leaf_elements (forest),
+        t8_forest_get_first_local_leaf_element_id (forest)
+      };
+    }
+
+    double local_partition_weight = 0.;
+    // t8_forest_get_num_local_trees & t8_forest_get_tree_num_leaf_elements assume a
+    // commited forest, which is not the case here yet...
+    for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
+      for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
+        local_partition_weight += weight_fcn (forest, ltreeid, ielm);
+      }
+    }
+
+    double local_partition_weight_offset = 0.;
+    sc_MPI_Scan (&local_partition_weight, &local_partition_weight_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
+    local_partition_weight_offset -= local_partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
+
+    double forest_weight = local_partition_weight_offset + local_partition_weight;
+    sc_MPI_Bcast (&forest_weight, 1, sc_MPI_DOUBLE, forest->mpisize - 1, forest->mpicomm);
+
+    return {
+        forest_weight,
+        local_partition_weight,
+        local_partition_weight_offset
+    };
+}
+
 /* Calculate the new element_offset for forest from
  * the element in forest->set_from using the provided weight function */
 static void
@@ -422,44 +461,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
   /* Initialize the shmem array */
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
 
-  // Compute the weight of all the partition-local elements as a whole
-  double const partition_weight = [&] () -> double {
-    if (weight_fcn == nullptr) {
-      return t8_forest_get_local_num_leaf_elements (forest);
-    }
-    double retval = 0.;
-    // t8_forest_get_num_local_trees & t8_forest_get_tree_num_leaf_elements assume a
-    // commited forest, which is not the case here yet...
-    for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
-      for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
-        retval += weight_fcn (forest, ltreeid, ielm);
-      }
-    }
-    return retval;
-  }();
-
-  // Compute the aggregated weight of all the partitions of lower rank
-  double const partition_weight_offset = [&] () -> double {
-    if (weight_fcn == nullptr) {
-      return t8_forest_get_first_local_leaf_element_id (forest);
-    }
-    double retval = 0.;
-    sc_MPI_Scan (&partition_weight, &retval, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
-    retval -= partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
-    return retval;
-  }();
-
-  // Compute the total weight of the forest and derive the equilibrium weight for each partition
-  double const total_forest_weight = [&] () -> double {
-    if (weight_fcn == nullptr) {
-      return t8_forest_get_global_num_leaf_elements (forest);
-    }
-    double retval = partition_weight_offset + partition_weight;
-    sc_MPI_Bcast (&retval, 1, sc_MPI_DOUBLE, mpisize - 1, forest->mpicomm);
-    return retval;
-  }();
-
-  double const target_weight = total_forest_weight / mpisize;
+  auto const [forest_weight, partition_weight, partition_weight_offset] = t8_forest_integrate_leaf_weights(forest, weight_fcn);
 
   // t8_forest_get_first_local_leaf_element_id too requires a commited forest I believe...
   t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id (forest);
@@ -472,13 +474,14 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
       t8_locidx_t current_elm_in_tree = 0;
       double accumulated_weight = partition_weight_offset;
       for (int i = 0; i < mpisize; i++) {
-        if (i * target_weight < partition_weight_offset
-            or i * target_weight > partition_weight_offset + partition_weight) {
+        double const target = forest_weight * i / mpisize;
+        if (target < partition_weight_offset
+            or target > partition_weight_offset + partition_weight) {
           continue;  // the new first element for the i-th partition is not here
         }
         // this while loop is a very ugly way to say "find the element where the partial sum
-        // of all the weights up to this point matches i*target_weight"
-        while (accumulated_weight < i * target_weight) {
+        // of all the weights up to this point matches the target"
+        while (accumulated_weight < target) {
           T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest));
           T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest, current_tree));
           if (weight_fcn == nullptr) {

From c7f3817614547de9efaf32cf5f2044658c527d5d Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 14 Oct 2025 15:19:18 +0200
Subject: [PATCH 09/35] use forest_from instead of forest

---
 src/t8_forest/t8_forest_partition.cxx | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 46e308429b..c287918f15 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -420,8 +420,6 @@ t8_forest_integrate_leaf_weights (t8_forest_t forest, weight_fcn_t *weight_fcn)
     }
 
     double local_partition_weight = 0.;
-    // t8_forest_get_num_local_trees & t8_forest_get_tree_num_leaf_elements assume a
-    // commited forest, which is not the case here yet...
     for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
       for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
         local_partition_weight += weight_fcn (forest, ltreeid, ielm);
@@ -461,10 +459,9 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
   /* Initialize the shmem array */
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
 
-  auto const [forest_weight, partition_weight, partition_weight_offset] = t8_forest_integrate_leaf_weights(forest, weight_fcn);
+  auto const [forest_weight, partition_weight, partition_weight_offset] = t8_forest_integrate_leaf_weights(forest_from, weight_fcn);
 
-  // t8_forest_get_first_local_leaf_element_id too requires a commited forest I believe...
-  t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id (forest);
+  t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id (forest_from);
 
   if (t8_shmem_array_start_writing (forest->element_offsets)) {
     if (forest_from->global_num_leaf_elements > 0) {
@@ -482,23 +479,23 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
         // this while loop is a very ugly way to say "find the element where the partial sum
         // of all the weights up to this point matches the target"
         while (accumulated_weight < target) {
-          T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest));
-          T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest, current_tree));
+          T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest_from));
+          T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest_from, current_tree));
           if (weight_fcn == nullptr) {
             accumulated_weight += 1;
           }
           else {
-            accumulated_weight += weight_fcn (forest, current_tree, current_elm_in_tree);
+            accumulated_weight += weight_fcn (forest_from, current_tree, current_elm_in_tree);
           }
           ++current_elm_in_tree;
-          if (current_elm_in_tree == t8_forest_get_tree_num_leaf_elements (forest, current_tree)) {
+          if (current_elm_in_tree == t8_forest_get_tree_num_leaf_elements (forest_from, current_tree)) {
             ++current_tree;
             current_elm_in_tree = 0;
           }
         }
         /* Calculate the first element index for the i-th process */
         t8_gloidx_t new_first_element_id
-          = partition_offset + t8_forest_get_tree_element_offset (forest, current_tree)  // same problem here
+          = partition_offset + t8_forest_get_tree_element_offset (forest_from, current_tree)
             + current_elm_in_tree;
         T8_ASSERT (0 <= new_first_element_id && new_first_element_id < forest_from->global_num_leaf_elements);
         element_offsets[i] = new_first_element_id;

From 3bfc468005f2ee8f05b07135b806f93bb5a09709 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 14 Oct 2025 15:20:49 +0200
Subject: [PATCH 10/35] indent

---
 src/t8_forest/t8_forest_partition.cxx | 53 ++++++++++++---------------
 1 file changed, 23 insertions(+), 30 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index c287918f15..36913683c6 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -406,38 +406,32 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
 // Computes the weight of the forest, the weight of the local partition, and the local offset (i.e. the weight
 // of all the partitions of lower rank combined)
 // If weight_fcn is null, all the elements are assumed to be of unit weight.
-static std::tuple< double, double, double >
+static std::tuple<double, double, double>
 t8_forest_integrate_leaf_weights (t8_forest_t forest, weight_fcn_t *weight_fcn)
 {
-    T8_ASSERT( t8_forest_is_committed(forest) );
-
-    if (weight_fcn == nullptr) {
-      return {
-        t8_forest_get_global_num_leaf_elements (forest),
-        t8_forest_get_local_num_leaf_elements (forest),
-        t8_forest_get_first_local_leaf_element_id (forest)
-      };
-    }
+  T8_ASSERT (t8_forest_is_committed (forest));
 
-    double local_partition_weight = 0.;
-    for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
-      for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
-        local_partition_weight += weight_fcn (forest, ltreeid, ielm);
-      }
+  if (weight_fcn == nullptr) {
+    return { t8_forest_get_global_num_leaf_elements (forest), t8_forest_get_local_num_leaf_elements (forest),
+             t8_forest_get_first_local_leaf_element_id (forest) };
+  }
+
+  double local_partition_weight = 0.;
+  for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
+    for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
+      local_partition_weight += weight_fcn (forest, ltreeid, ielm);
     }
+  }
 
-    double local_partition_weight_offset = 0.;
-    sc_MPI_Scan (&local_partition_weight, &local_partition_weight_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
-    local_partition_weight_offset -= local_partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
+  double local_partition_weight_offset = 0.;
+  sc_MPI_Scan (&local_partition_weight, &local_partition_weight_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
+  local_partition_weight_offset
+    -= local_partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
 
-    double forest_weight = local_partition_weight_offset + local_partition_weight;
-    sc_MPI_Bcast (&forest_weight, 1, sc_MPI_DOUBLE, forest->mpisize - 1, forest->mpicomm);
+  double forest_weight = local_partition_weight_offset + local_partition_weight;
+  sc_MPI_Bcast (&forest_weight, 1, sc_MPI_DOUBLE, forest->mpisize - 1, forest->mpicomm);
 
-    return {
-        forest_weight,
-        local_partition_weight,
-        local_partition_weight_offset
-    };
+  return { forest_weight, local_partition_weight, local_partition_weight_offset };
 }
 
 /* Calculate the new element_offset for forest from
@@ -459,7 +453,8 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
   /* Initialize the shmem array */
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
 
-  auto const [forest_weight, partition_weight, partition_weight_offset] = t8_forest_integrate_leaf_weights(forest_from, weight_fcn);
+  auto const [forest_weight, partition_weight, partition_weight_offset]
+    = t8_forest_integrate_leaf_weights (forest_from, weight_fcn);
 
   t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id (forest_from);
 
@@ -472,8 +467,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
       double accumulated_weight = partition_weight_offset;
       for (int i = 0; i < mpisize; i++) {
         double const target = forest_weight * i / mpisize;
-        if (target < partition_weight_offset
-            or target > partition_weight_offset + partition_weight) {
+        if (target < partition_weight_offset or target > partition_weight_offset + partition_weight) {
           continue;  // the new first element for the i-th partition is not here
         }
         // this while loop is a very ugly way to say "find the element where the partial sum
@@ -495,8 +489,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
         }
         /* Calculate the first element index for the i-th process */
         t8_gloidx_t new_first_element_id
-          = partition_offset + t8_forest_get_tree_element_offset (forest_from, current_tree)
-            + current_elm_in_tree;
+          = partition_offset + t8_forest_get_tree_element_offset (forest_from, current_tree) + current_elm_in_tree;
         T8_ASSERT (0 <= new_first_element_id && new_first_element_id < forest_from->global_num_leaf_elements);
         element_offsets[i] = new_first_element_id;
       }

From 56b51b1467b19b282064177961e2a35aaeb17ed8 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Wed, 15 Oct 2025 12:39:34 +0200
Subject: [PATCH 11/35] restore optimized path for uniform weights

---
 src/t8_forest/t8_forest_partition.cxx | 53 ++++++++++++++-------------
 1 file changed, 28 insertions(+), 25 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 36913683c6..da302a2727 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -462,38 +462,41 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
     if (forest_from->global_num_leaf_elements > 0) {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
 
-      t8_locidx_t current_tree = 0;
-      t8_locidx_t current_elm_in_tree = 0;
-      double accumulated_weight = partition_weight_offset;
-      for (int i = 0; i < mpisize; i++) {
-        double const target = forest_weight * i / mpisize;
-        if (target < partition_weight_offset or target > partition_weight_offset + partition_weight) {
-          continue;  // the new first element for the i-th partition is not here
+      if (weight_fcn == nullptr) {  // in this case forest_weight = global_num_leaf_elements
+        for (int i = 0; i < mpisize; i++) {
+          element_offsets[i] = static_cast< t8_gloidx_t >( forest_weight * i / mpisize );
         }
-        // this while loop is a very ugly way to say "find the element where the partial sum
-        // of all the weights up to this point matches the target"
-        while (accumulated_weight < target) {
-          T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest_from));
-          T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest_from, current_tree));
-          if (weight_fcn == nullptr) {
-            accumulated_weight += 1;
+      }
+      else {
+        t8_locidx_t current_tree = 0;
+        t8_locidx_t current_elm_in_tree = 0;
+        double accumulated_weight = partition_weight_offset;
+        for (int i = 0; i < mpisize; i++) {
+          double const target = forest_weight * i / mpisize;
+          if (target < partition_weight_offset or target > partition_weight_offset + partition_weight) {
+            continue;  // the new first element for the i-th partition is not here
           }
-          else {
+          // this while loop is a very ugly way to say "find the element where the partial sum
+          // of all the weights up to this point matches the target"
+          while (accumulated_weight < target) {
+            T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest_from));
+            T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest_from, current_tree));
             accumulated_weight += weight_fcn (forest_from, current_tree, current_elm_in_tree);
+            ++current_elm_in_tree;
+            if (current_elm_in_tree == t8_forest_get_tree_num_leaf_elements (forest_from, current_tree)) {
+              ++current_tree;
+              current_elm_in_tree = 0;
+            }
           }
-          ++current_elm_in_tree;
-          if (current_elm_in_tree == t8_forest_get_tree_num_leaf_elements (forest_from, current_tree)) {
-            ++current_tree;
-            current_elm_in_tree = 0;
-          }
+          element_offsets[i]
+            = partition_offset + t8_forest_get_tree_element_offset (forest_from, current_tree) + current_elm_in_tree;
         }
-        /* Calculate the first element index for the i-th process */
-        t8_gloidx_t new_first_element_id
-          = partition_offset + t8_forest_get_tree_element_offset (forest_from, current_tree) + current_elm_in_tree;
-        T8_ASSERT (0 <= new_first_element_id && new_first_element_id < forest_from->global_num_leaf_elements);
-        element_offsets[i] = new_first_element_id;
       }
       element_offsets[forest->mpisize] = forest->global_num_leaf_elements;
+
+      for (int i = 0; i < mpisize; i++) {
+        T8_ASSERT (0 <= element_offsets[i] && element_offsets[i] < forest_from->global_num_leaf_elements);
+      }
     }
     else {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);

From c3ef88d2a035498290a3d8d93f0208ad87f4dce5 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Wed, 15 Oct 2025 13:20:08 +0200
Subject: [PATCH 12/35] remove obfuscated while loop

---
 src/t8_forest/t8_forest_partition.cxx | 31 ++++++++++-----------------
 1 file changed, 11 insertions(+), 20 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index da302a2727..2bffd37132 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -464,32 +464,23 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight
 
       if (weight_fcn == nullptr) {  // in this case forest_weight = global_num_leaf_elements
         for (int i = 0; i < mpisize; i++) {
-          element_offsets[i] = static_cast< t8_gloidx_t >( forest_weight * i / mpisize );
+          element_offsets[i] = static_cast<t8_gloidx_t> (forest_weight * i / mpisize);
         }
       }
       else {
-        t8_locidx_t current_tree = 0;
-        t8_locidx_t current_elm_in_tree = 0;
         double accumulated_weight = partition_weight_offset;
-        for (int i = 0; i < mpisize; i++) {
-          double const target = forest_weight * i / mpisize;
-          if (target < partition_weight_offset or target > partition_weight_offset + partition_weight) {
-            continue;  // the new first element for the i-th partition is not here
-          }
-          // this while loop is a very ugly way to say "find the element where the partial sum
-          // of all the weights up to this point matches the target"
-          while (accumulated_weight < target) {
-            T8_ASSERT (current_tree < t8_forest_get_num_local_trees (forest_from));
-            T8_ASSERT (current_elm_in_tree < t8_forest_get_tree_num_leaf_elements (forest_from, current_tree));
-            accumulated_weight += weight_fcn (forest_from, current_tree, current_elm_in_tree);
-            ++current_elm_in_tree;
-            if (current_elm_in_tree == t8_forest_get_tree_num_leaf_elements (forest_from, current_tree)) {
-              ++current_tree;
-              current_elm_in_tree = 0;
+        t8_gloidx_t global_elm_idx = partition_offset;
+        int i = std::ceil (partition_weight_offset * mpisize / forest_weight);
+
+        for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
+          for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
+            accumulated_weight += weight_fcn (forest, ltreeid, ielm);
+            ++global_elm_idx;
+            if (accumulated_weight > forest_weight * i / mpisize) {
+              element_offsets[i] = global_elm_idx;
+              ++i;
             }
           }
-          element_offsets[i]
-            = partition_offset + t8_forest_get_tree_element_offset (forest_from, current_tree) + current_elm_in_tree;
         }
       }
       element_offsets[forest->mpisize] = forest->global_num_leaf_elements;

From 0f3f711b2dd9b729f15fcdacdb196e113b5d1145 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 28 Oct 2025 11:08:02 +0100
Subject: [PATCH 13/35] expose the weight function through
 t8_forest_set_partition

---
 benchmarks/t8_time_forest_partition.cxx       |  2 +-
 benchmarks/t8_time_prism_adapt.cxx            |  4 +-
 example/advect/t8_advection.cxx               |  2 +-
 example/forest/t8_test_face_iterate.cxx       |  4 +-
 example/forest/t8_test_ghost.cxx              |  4 +-
 .../forest/t8_test_ghost_large_level_diff.cxx |  4 +-
 example/remove/t8_example_empty_trees.cxx     |  4 +-
 src/t8_forest/t8_forest.cxx                   | 39 +++++++-------
 src/t8_forest/t8_forest_balance.cxx           |  2 +-
 src/t8_forest/t8_forest_general.h             | 52 +++++++++++--------
 src/t8_forest/t8_forest_partition.cxx         |  6 +--
 src/t8_forest/t8_forest_partition.h           |  8 +--
 src/t8_forest/t8_forest_types.h               | 13 ++---
 test/t8_forest/t8_gtest_forest_commit.cxx     |  6 +--
 test/t8_forest/t8_gtest_partition_data.cxx    |  2 +-
 .../t8_gtest_empty_global_tree.cxx            | 12 ++---
 .../t8_gtest_empty_local_tree.cxx             | 24 ++++-----
 .../t8_gtest_iterate_replace.cxx              |  6 +--
 .../t8_gtest_element_ref_coords.cxx           |  2 +-
 .../features/t8_features_curved_meshes.cxx    | 32 ++++++------
 .../t8_step4_partition_balance_ghost.cxx      | 40 +++++++-------
 tutorials/general/t8_step5_element_data.cxx   | 18 +++----
 .../t8_step5_element_data_c_interface.c       | 18 +++----
 tutorials/general/t8_step6_stencil.cxx        | 10 ++--
 24 files changed, 159 insertions(+), 155 deletions(-)

diff --git a/benchmarks/t8_time_forest_partition.cxx b/benchmarks/t8_time_forest_partition.cxx
index fa28d0caf2..558ecbd2f7 100644
--- a/benchmarks/t8_time_forest_partition.cxx
+++ b/benchmarks/t8_time_forest_partition.cxx
@@ -179,7 +179,7 @@ t8_time_forest_cmesh_mshfile (t8_cmesh_t cmesh, const char *vtu_prefix, sc_MPI_C
     /* partition the adapted forest */
     t8_forest_init (&forest_partition);
     /* partition the adapted forest */
-    t8_forest_set_partition (forest_partition, forest_adapt, 0);
+    t8_forest_set_partition (forest_partition, forest_adapt, 0, nullptr);
 
     /* If desired, create ghost elements and balance */
     t8_forest_set_profiling (forest_partition, 1);
diff --git a/benchmarks/t8_time_prism_adapt.cxx b/benchmarks/t8_time_prism_adapt.cxx
index 2d7784fe07..3b32c5ec4d 100644
--- a/benchmarks/t8_time_prism_adapt.cxx
+++ b/benchmarks/t8_time_prism_adapt.cxx
@@ -128,7 +128,7 @@ t8_time_refine (int start_level, int end_level, [[maybe_unused]] int create_fore
     }
     forest_partition = forest_adapt;
     /* partition the adapted forest */
-    t8_forest_set_partition (forest_partition, NULL, 0);
+    t8_forest_set_partition (forest_partition, NULL, 0, nullptr);
     /* enable profiling for the partitioned forest */
     t8_forest_set_profiling (forest_partition, 1);
     /* if desired do balance */
@@ -181,7 +181,7 @@ main (int argc, char **argv)
 
   if (sreturnA > BUFSIZ || sreturnB > BUFSIZ) {
     /* The usage string or help message was truncated */
-    /* Note: gcc >= 7.1 prints a warning if we 
+    /* Note: gcc >= 7.1 prints a warning if we
      * do not check the return value of snprintf. */
     t8_debugf ("Warning: Truncated usage string and help message to '%s' and '%s'\n", usage, help);
   }
diff --git a/example/advect/t8_advection.cxx b/example/advect/t8_advection.cxx
index 7ae4136716..5aeccd6622 100644
--- a/example/advect/t8_advection.cxx
+++ b/example/advect/t8_advection.cxx
@@ -784,7 +784,7 @@ t8_advect_problem_partition (t8_advect_problem_t *problem, int measure_time)
   /* Enable profiling to measure runtime */
   t8_forest_set_profiling (forest_partition, 1);
   /* Partition the forest and create ghosts */
-  t8_forest_set_partition (forest_partition, problem->forest, 0);
+  t8_forest_set_partition (forest_partition, problem->forest, 0, nullptr);
   t8_forest_set_ghost (forest_partition, 1, T8_GHOST_FACES);
   t8_forest_commit (forest_partition);
 
diff --git a/example/forest/t8_test_face_iterate.cxx b/example/forest/t8_test_face_iterate.cxx
index 63b4ed915b..15631dc8c0 100644
--- a/example/forest/t8_test_face_iterate.cxx
+++ b/example/forest/t8_test_face_iterate.cxx
@@ -138,7 +138,7 @@ t8_test_fiterate_refine_and_partition (t8_cmesh_t cmesh, int level, sc_MPI_Comm
 
   /* partition the adapted forest */
   t8_forest_init (&forest_partition);
-  t8_forest_set_partition (forest_partition, forest_adapt, 0);
+  t8_forest_set_partition (forest_partition, forest_adapt, 0, nullptr);
   t8_forest_commit (forest_partition);
   t8_debugf ("Created ghost structure with %li ghost elements.\n", (long) t8_forest_get_num_ghosts (forest_partition));
   if (!no_vtk) {
@@ -212,7 +212,7 @@ main (int argc, char **argv)
 
   if (sreturnA > BUFSIZ || sreturnB > BUFSIZ) {
     /* The usage string or help message was truncated */
-    /* Note: gcc >= 7.1 prints a warning if we 
+    /* Note: gcc >= 7.1 prints a warning if we
      * do not check the return value of snprintf. */
     t8_debugf ("Warning: Truncated usage string and help message to '%s' and '%s'\n", usage, help);
   }
diff --git a/example/forest/t8_test_ghost.cxx b/example/forest/t8_test_ghost.cxx
index f941da62be..10c4a7976b 100644
--- a/example/forest/t8_test_ghost.cxx
+++ b/example/forest/t8_test_ghost.cxx
@@ -189,7 +189,7 @@ t8_test_ghost_refine_and_partition (t8_cmesh_t cmesh, const int level, sc_MPI_Co
   }
 
   /* Set the forest for partitioning */
-  t8_forest_set_partition (forest_ghost, forest, 0);
+  t8_forest_set_partition (forest_ghost, forest, 0, nullptr);
   /* Activate ghost creation */
   t8_forest_set_ghost_ext (forest_ghost, 1, T8_GHOST_FACES, ghost_version);
   /* Activate timers */
@@ -303,7 +303,7 @@ main (int argc, char **argv)
 
   if (sreturnA > BUFSIZ || sreturnB > BUFSIZ) {
     /* The usage string or help message was truncated */
-    /* Note: gcc >= 7.1 prints a warning if we 
+    /* Note: gcc >= 7.1 prints a warning if we
      * do not check the return value of snprintf. */
     t8_debugf ("Warning: Truncated usage string and help message to '%s' and '%s'\n", usage, help);
   }
diff --git a/example/forest/t8_test_ghost_large_level_diff.cxx b/example/forest/t8_test_ghost_large_level_diff.cxx
index 56911715c3..2a68556367 100644
--- a/example/forest/t8_test_ghost_large_level_diff.cxx
+++ b/example/forest/t8_test_ghost_large_level_diff.cxx
@@ -179,7 +179,7 @@ t8_ghost_large_level_diff (const char *prefix, int dim, int level, int refine, i
 
   /* Partition */
   t8_forest_init (&forest_partition);
-  t8_forest_set_partition (forest_partition, forest_adapt, 0);
+  t8_forest_set_partition (forest_partition, forest_adapt, 0, nullptr);
   t8_forest_set_ghost_ext (forest_partition, 1, T8_GHOST_FACES, 3);
   t8_forest_set_profiling (forest_partition, 1);
   t8_forest_commit (forest_partition);
@@ -219,7 +219,7 @@ main (int argc, char *argv[])
 
   if (sreturn >= BUFSIZ) {
     /* The help message was truncated */
-    /* Note: gcc >= 7.1 prints a warning if we 
+    /* Note: gcc >= 7.1 prints a warning if we
      * do not check the return value of snprintf. */
     t8_debugf ("Warning: Truncated help message to '%s'\n", help);
   }
diff --git a/example/remove/t8_example_empty_trees.cxx b/example/remove/t8_example_empty_trees.cxx
index e2a95c3fd7..1c5e487b73 100644
--- a/example/remove/t8_example_empty_trees.cxx
+++ b/example/remove/t8_example_empty_trees.cxx
@@ -66,7 +66,7 @@ t8_strip_of_quads (t8_gloidx_t num_trees, t8_gloidx_t empty_tree, const char **v
   t8_forest_t forest_adapt;
   t8_forest_init (&forest_adapt);
   t8_forest_set_adapt (forest_adapt, forest, t8_adapt_remove, 0);
-  t8_forest_set_partition (forest_adapt, NULL, 0);
+  t8_forest_set_partition (forest_adapt, NULL, 0, nullptr);
   t8_forest_set_user_data (forest_adapt, &empty_tree);
   t8_forest_commit (forest_adapt);
 
@@ -115,7 +115,7 @@ main (int argc, char **argv)
 
   if (sreturnA > BUFSIZ || sreturnB > BUFSIZ) {
     /* The usage string or help message was truncated */
-    /* Note: gcc >= 7.1 prints a warning if we 
+    /* Note: gcc >= 7.1 prints a warning if we
      * do not check the return value of snprintf. */
     t8_debugf ("Warning: Truncated usage string and help message to '%s' and '%s'\n", usage, help);
   }
diff --git a/src/t8_forest/t8_forest.cxx b/src/t8_forest/t8_forest.cxx
index b607bb1fef..7fe3d920b1 100644
--- a/src/t8_forest/t8_forest.cxx
+++ b/src/t8_forest/t8_forest.cxx
@@ -79,7 +79,7 @@ t8_forest_is_incomplete_family (const t8_forest_t forest, const t8_locidx_t ltre
   scheme->element_new (tree_class, 1, &element_parent_current);
   scheme->element_new (tree_class, 1, &element_compare);
 
-  /* We first assume that we have an (in)complete family with the size of array elements. 
+  /* We first assume that we have an (in)complete family with the size of array elements.
    * In the following we try to disprove this. */
   int family_size = elements_size;
 
@@ -88,9 +88,9 @@ t8_forest_is_incomplete_family (const t8_forest_t forest, const t8_locidx_t ltre
   const int child_id_current = scheme->element_get_child_id (tree_class, elements[0]);
   scheme->element_get_parent (tree_class, elements[0], element_parent_current);
 
-  /* Elements of the current family could already be passed, so that 
+  /* Elements of the current family could already be passed, so that
    * the element/family currently under consideration can no longer be coarsened.
-   * Also, there may be successors of a hypothetical previous family member 
+   * Also, there may be successors of a hypothetical previous family member
    * that would be overlapped after coarsening.
    * */
   if (child_id_current > 0 && el_considered > 0) {
@@ -132,12 +132,12 @@ t8_forest_is_incomplete_family (const t8_forest_t forest, const t8_locidx_t ltre
   T8_ASSERT (family_size > 0);
   T8_ASSERT (family_size >= 0 && family_size <= elements_size);
 
-  /* There may be successors of a hypothetical later family member (with index 
+  /* There may be successors of a hypothetical later family member (with index
    * family_size in this family) that would be overlapped after coarsening. */
   if (family_size < elements_size) {
     /* Get level of element after last element of current possible family */
     const int level = scheme->element_get_level (tree_class, elements[family_size]);
-    /* Only elements with higher level then level of current element, can get 
+    /* Only elements with higher level then level of current element, can get
      * potentially be overlapped. */
     if (level > level_current) {
       /* Compare ancestors */
@@ -160,7 +160,7 @@ t8_forest_is_incomplete_family (const t8_forest_t forest, const t8_locidx_t ltre
   const int num_siblings = scheme->element_get_num_siblings (tree_class, elements[0]);
   T8_ASSERT (family_size <= num_siblings);
   /* If the first/last element at a process boundary is not the first/last
-   * element of a possible family, we are not guaranteed to consider all 
+   * element of a possible family, we are not guaranteed to consider all
    * family members.*/
   if (el_considered == 0 && child_id_current > 0 && ltree_id == 0 && forest->mpirank > 0) {
     return 0;
@@ -296,9 +296,9 @@ t8_forest_no_overlap ([[maybe_unused]] t8_forest_t forest)
        * More detailed:
        * Let e_a and e_b be two elements.
        * If the level of e_a is equal to the level of the nca of e_a and e_b,
-       * then e_b is a descendant of e_a. 
+       * then e_b is a descendant of e_a.
        * If the level of e_b is equal to the level of the nca of e_a and e_b,
-       * then e_a is a descendant of e_b. 
+       * then e_a is a descendant of e_b.
        * Thus e_a and e_b overlap in both cases.
        * Note: If e_a equals e_b, e_a is the descendant of e_b and vice versa.
        * */
@@ -1297,9 +1297,9 @@ t8_forest_tree_shared ([[maybe_unused]] t8_forest_t forest, [[maybe_unused]] int
     else {
       SC_ABORT ("For incomplete trees the method t8_forest_last_tree_shared aka "
                 "t8_forest_tree_shared(forest, 1) is not implemented.\n");
-      /* TODO: If last_local_tree is 0 of the current process and it gets 0 as the 
-       * first_local_tree of the bigger process, then it cannot be said whether 
-       * the tree with id 0 is shared or not, since the bigger process could also 
+      /* TODO: If last_local_tree is 0 of the current process and it gets 0 as the
+       * first_local_tree of the bigger process, then it cannot be said whether
+       * the tree with id 0 is shared or not, since the bigger process could also
        * carry an empty forest. */
     }
     /* If global_neighbour_tree_idx == forest->first_local_tree tree is shared */
@@ -1964,7 +1964,7 @@ t8_forest_element_is_leaf (const t8_forest_t forest, const t8_element_t *element
   T8_ASSERT (t8_forest_is_committed (forest));
   T8_ASSERT (t8_forest_tree_is_local (forest, local_tree));
 
-  /* We get the array of the tree's elements and then search in the array of elements for our 
+  /* We get the array of the tree's elements and then search in the array of elements for our
    * element candidate. */
   /* Get the array */
   const t8_element_array_t *elements = t8_forest_get_tree_leaf_element_array (forest, local_tree);
@@ -1986,7 +1986,7 @@ t8_forest_element_is_leaf (const t8_forest_t forest, const t8_element_t *element
     /* The element was not found. */
     return 0;
   }
-  /* An element was found but it may not be the candidate element. 
+  /* An element was found but it may not be the candidate element.
    * To identify whether the element was found, we compare these two. */
   const t8_element_t *check_element = t8_element_array_index_locidx (elements, search_result);
   T8_ASSERT (check_element != NULL);
@@ -2327,7 +2327,7 @@ t8_forest_element_find_owner_old (t8_forest_t forest, t8_gloidx_t gtreeid, t8_el
     return proc;
   }
   else {
-    /* Get the next owning process. Its first descendant is in fact an element of the tree. 
+    /* Get the next owning process. Its first descendant is in fact an element of the tree.
      * If it is bigger than the descendant we look for, then proc is the owning process of element. */
     proc_next = *(int *) sc_array_index (owners_of_tree, 1);
     if (*(t8_linearidx_t *) t8_shmem_array_index (forest->global_first_desc, (size_t) proc_next)
@@ -2814,7 +2814,7 @@ t8_forest_set_copy (t8_forest_t forest, const t8_forest_t set_from)
 }
 
 void
-t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening)
+t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening, t8_weight_fcn_t* weight_callback)
 {
   T8_ASSERT (forest != NULL);
   T8_ASSERT (forest->rc.refcount > 0);
@@ -2824,6 +2824,7 @@ t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set
   T8_ASSERT (forest->scheme == NULL);
 
   forest->set_for_coarsening = set_for_coarsening;
+  forest->weight_function = weight_callback;
 
   if (set_from != NULL) {
     /* If set_from = NULL, we assume a previous forest_from was set */
@@ -2972,9 +2973,9 @@ t8_forest_comm_global_num_leaf_elements (t8_forest_t forest)
  * Check if any tree in a forest refines irregularly.
  * An irregular refining tree is a tree with an element that does not
  * refine into 2^dim children. For example the default implementation
- * of pyramids. 
+ * of pyramids.
  * \note This function is MPI collective
- * 
+ *
  * \param[in] forest    The forest to check
  * \return          Non-zero if any tree refines irregular
  */
@@ -3164,7 +3165,7 @@ t8_forest_commit (t8_forest_t forest)
           /* forest_partition should not change ownership of forest->set_from */
           t8_forest_ref (forest->set_from);
         }
-        t8_forest_set_partition (forest_partition, forest->set_from, forest->set_for_coarsening);
+        t8_forest_set_partition (forest_partition, forest->set_from, forest->set_for_coarsening, forest->weight_function);
         /* activate profiling, if this forest has profiling */
         t8_forest_set_profiling (forest_partition, forest->profile != NULL);
         /* Commit the partitioned forest */
@@ -3185,7 +3186,7 @@ t8_forest_commit (t8_forest_t forest)
         /* Initialize the trees array of the forest */
         forest->trees = sc_array_new (sizeof (t8_tree_struct_t));
         /* partition the forest */
-        t8_forest_partition (forest);
+        t8_forest_partition (forest, forest->weight_function);
       }
     }
     if (forest->from_method & T8_FOREST_FROM_BALANCE) {
diff --git a/src/t8_forest/t8_forest_balance.cxx b/src/t8_forest/t8_forest_balance.cxx
index 7381628049..50086a03ab 100644
--- a/src/t8_forest/t8_forest_balance.cxx
+++ b/src/t8_forest/t8_forest_balance.cxx
@@ -245,7 +245,7 @@ t8_forest_balance (t8_forest_t forest, int repartition)
       t8_forest_init (&forest_partition);
       /* Update the maximum occurring level */
       forest_partition->maxlevel_existing = forest_temp->maxlevel_existing;
-      t8_forest_set_partition (forest_partition, forest_temp, 0);
+      t8_forest_set_partition (forest_partition, forest_temp, 0, nullptr);
       t8_forest_set_ghost (forest_partition, 1, T8_GHOST_FACES);
       /* If profiling is enabled, measure partition runtimes */
       if (forest->profile != NULL) {
diff --git a/src/t8_forest/t8_forest_general.h b/src/t8_forest/t8_forest_general.h
index 3613430cf1..49bae28651 100644
--- a/src/t8_forest/t8_forest_general.h
+++ b/src/t8_forest/t8_forest_general.h
@@ -48,11 +48,18 @@ typedef enum {
   T8_GHOST_VERTICES  /**< Consider all vertex (codimension 3) and edge and face neighbors. */
 } t8_ghost_type_t;
 
-/** This typedef is needed as a helper construct to 
+/** This typedef is needed as a helper construct to
  * properly be able to define a function that returns
  * a pointer to a void fun(void) function. \see t8_forest_get_user_function.
  */
 typedef void (*t8_generic_function_pointer) (void);
+
+/**
+ * The prototype a weight function for the partition algorithm.
+ * The function should be pure, and return a positive weight given a forest, a local tree index and an element index within the local tree
+ */
+typedef double (t8_weight_fcn_t) (t8_forest_t, t8_locidx_t, t8_locidx_t);
+
 T8_EXTERN_C_BEGIN ();
 
 /** Callback function prototype to replace one set of elements with another.
@@ -77,12 +84,12 @@ T8_EXTERN_C_BEGIN ();
  * \param [in] first_incoming  The tree local index of the first incoming element.
  *                             0 <= first_incom < new_which_tree->num_elements
  *
- * If an element is being refined, \a refine and \a num_outgoing will be 1 and 
+ * If an element is being refined, \a refine and \a num_outgoing will be 1 and
  * \a num_incoming will be the number of children.
- * If a family is being coarsened, \a refine will be -1, \a num_outgoing will be 
- * the number of family members and \a num_incoming will be 1. 
- * If an element is being removed, \a refine and \a num_outgoing will be 1 and 
- * \a num_incoming will be 0. 
+ * If a family is being coarsened, \a refine will be -1, \a num_outgoing will be
+ * the number of family members and \a num_incoming will be 1.
+ * If an element is being removed, \a refine and \a num_outgoing will be 1 and
+ * \a num_incoming will be 0.
  * Else \a refine will be 0 and \a num_outgoing and \a num_incoming will both be 1.
  * \see t8_forest_iterate_replace
  */
@@ -96,7 +103,7 @@ typedef void (*t8_forest_replace_t) (t8_forest_t forest_old, t8_forest_t forest_
  * form a family and we decide whether this family should be coarsened
  * or only the first element should be refined.
  * Otherwise \a is_family must equal zero and we consider the first entry
- * of the element array for refinement. 
+ * of the element array for refinement.
  * Entries of the element array beyond the first \a num_elements are undefined.
  * \param [in] forest       The forest to which the new elements belong.
  * \param [in] forest_from  The forest that is adapted.
@@ -121,7 +128,7 @@ typedef int (*t8_forest_adapt_t) (t8_forest_t forest, t8_forest_t forest_from, t
 
 /** Create a new forest with reference count one.
  * This forest needs to be specialized with the t8_forest_set_* calls.
- * Currently it is mandatory to either call the functions \see t8_forest_set_mpicomm, 
+ * Currently it is mandatory to either call the functions \see t8_forest_set_mpicomm,
  * \ref t8_forest_set_cmesh, and \ref t8_forest_set_scheme,
  * or to call one of \ref t8_forest_set_copy, \ref t8_forest_set_adapt, or
  * \ref t8_forest_set_partition.  It is illegal to mix these calls, or to
@@ -159,7 +166,7 @@ t8_forest_is_committed (t8_forest_t forest);
  * \param [in] forest   The forest to consider.
  * \return              True if \a forest has no elements which are inside each other.
  * \note This function is collective, but only checks local overlapping on each process.
- * \see t8_forest_partition_test_boundary_element if you also want to test for 
+ * \see t8_forest_partition_test_boundary_element if you also want to test for
  * global overlap across the process boundaries.
  */
 int
@@ -313,6 +320,7 @@ t8_forest_get_user_function (const t8_forest_t forest);
  * \param [in]      set_for_coarsening CURRENTLY DISABLED. If true, then the partitions
  *                          are choose such that coarsening an element once is a process local
  *                          operation.
+ * \param [in]      weight_callback A callback function defining element weights for the partitioning. If null, all the elements are assumed to have the same weight. Must be free of side-effects.
  * \note This setting can be combined with \ref t8_forest_set_adapt and \ref
  * t8_forest_set_balance. The order in which these operations are executed is always
  * 1) Adapt 2) Partition 3) Balance.
@@ -322,7 +330,7 @@ t8_forest_get_user_function (const t8_forest_t forest);
  * this setting.
  */
 void
-t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening);
+t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening, t8_weight_fcn_t *weight_callback);
 
 /** Set a source forest to be balanced during commit.
  * A forest is said to be balanced if each element has face neighbors of level
@@ -379,8 +387,8 @@ t8_forest_set_ghost_ext (t8_forest_t forest, int do_ghost, t8_ghost_type_t ghost
 
 /**
  *  Use assertions and document that the forest_set (..., from) and
- *  set_load are mutually exclusive. 
- * 
+ *  set_load are mutually exclusive.
+ *
  *  TODO: Unused function -> remove?
  */
 void
@@ -452,7 +460,7 @@ t8_forest_get_eclass (const t8_forest_t forest, const t8_locidx_t ltreeid);
 
 /**
  * Check whether a given tree id belongs to a local tree in a forest.
- * 
+ *
  * \param [in]    forest The forest.
  * \param [in]    local_tree A tree id.
  * \return True if and only if the id \a local_tree belongs to a local tree of \a forest.
@@ -520,7 +528,7 @@ t8_forest_get_coarse_tree (t8_forest_t forest, t8_locidx_t ltreeid);
 
 /**
  * Query whether a given element is a leaf in a forest.
- * 
+ *
  * \param [in]  forest    The forest.
  * \param [in]  element   An element of a local tree in \a forest.
  * \param [in]  local_tree A local tree id of \a forest.
@@ -583,7 +591,7 @@ t8_forest_leaf_face_neighbors (t8_forest_t forest, t8_locidx_t ltreeid, const t8
                                t8_element_t **pneighbor_leaves[], int face, int *dual_faces[], int *num_neighbors,
                                t8_locidx_t **pelement_indices, t8_eclass_t *pneigh_eclass, int forest_is_balanced);
 
-/** Like \ref t8_forest_leaf_face_neighbors but also provides information about the global neighbors and the orientation. 
+/** Like \ref t8_forest_leaf_face_neighbors but also provides information about the global neighbors and the orientation.
  * \param [in]    forest  The forest. Must have a valid ghost layer.
  * \param [in]    ltreeid A local tree id.
  * \param [in]    leaf    A leaf in tree \a ltreeid of \a forest.
@@ -601,8 +609,8 @@ t8_forest_leaf_face_neighbors (t8_forest_t forest, t8_locidx_t ltreeid, const t8
  * \param [in]    forest_is_balanced True if we know that \a forest is balanced, false
  *                        otherwise.
  * \param [out]   gneigh_tree  The global tree IDs of the neighbor trees.
- * \param [out]   orientation  If not NULL on input, the face orientation is computed and stored here. 
- *                                         Thus, if the face connection is an inter-tree connection the orientation of the tree-to-tree connection is stored. 
+ * \param [out]   orientation  If not NULL on input, the face orientation is computed and stored here.
+ *                                         Thus, if the face connection is an inter-tree connection the orientation of the tree-to-tree connection is stored.
  *                                         Otherwise, the value 0 is stored.
  * All other parameters and behavior are identical to \ref t8_forest_leaf_face_neighbors.
  * \note If there are no face neighbors, then *neighbor_leaves = NULL, num_neighbors = 0,
@@ -849,7 +857,7 @@ t8_forest_element_face_neighbor (t8_forest_t forest, t8_locidx_t ltreeid, const
 
 /**
  * TODO: Can be removed since it is unused.
- * 
+ *
  * \param[in] forest The forest.
  */
 void
@@ -857,16 +865,16 @@ t8_forest_iterate (t8_forest_t forest);
 
 /** Query whether a batch of points lies inside an element. For bilinearly interpolated elements.
  * \note For 2D quadrilateral elements this function is only an approximation. It is correct
- *  if the four vertices lie in the same plane, but it may produce only approximate results if 
+ *  if the four vertices lie in the same plane, but it may produce only approximate results if
  *  the vertices do not lie in the same plane.
  * \param [in]      forest      The forest.
  * \param [in]      ltreeid     The forest local id of the tree in which the element is.
  * \param [in]      element     The element.
  * \param [in]      points      3-dimensional coordinates of the points to check
  * \param [in]      num_points  The number of points to check
- * \param [in, out] is_inside   An array of length \a num_points, filled with 0/1 on output. True (non-zero) if a \a point 
- *                              lies within an \a element, false otherwise. The return value is also true if the point 
- *                              lies on the element boundary. Thus, this function may return true for different leaf 
+ * \param [in, out] is_inside   An array of length \a num_points, filled with 0/1 on output. True (non-zero) if a \a point
+ *                              lies within an \a element, false otherwise. The return value is also true if the point
+ *                              lies on the element boundary. Thus, this function may return true for different leaf
  *                              elements, if they are neighbors and the point lies on the common boundary.
  * \param [in]      tolerance   Tolerance that we allow the point to not exactly match the element.
  *                              If this value is larger we detect more points.
diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 2bffd37132..429f9bc8d6 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -407,7 +407,7 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
 // of all the partitions of lower rank combined)
 // If weight_fcn is null, all the elements are assumed to be of unit weight.
 static std::tuple<double, double, double>
-t8_forest_integrate_leaf_weights (t8_forest_t forest, weight_fcn_t *weight_fcn)
+t8_forest_integrate_leaf_weights (t8_forest_t forest, t8_weight_fcn_t *weight_fcn)
 {
   T8_ASSERT (t8_forest_is_committed (forest));
 
@@ -437,7 +437,7 @@ t8_forest_integrate_leaf_weights (t8_forest_t forest, weight_fcn_t *weight_fcn)
 /* Calculate the new element_offset for forest from
  * the element in forest->set_from using the provided weight function */
 static void
-t8_forest_partition_compute_new_offset (t8_forest_t forest, weight_fcn_t *weight_fcn)
+t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *weight_fcn)
 {
   T8_ASSERT (t8_forest_is_initialized (forest));
   T8_ASSERT (forest->set_from != NULL);
@@ -1219,7 +1219,7 @@ t8_forest_partition_given (t8_forest_t forest, const int send_data, const sc_arr
  * Currently the elements are distributed evenly (each element has the same weight).
  */
 void
-t8_forest_partition (t8_forest_t forest, weight_fcn_t *weight_callback)
+t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback)
 {
   t8_forest_t forest_from;
   int create_offset_from = 0;
diff --git a/src/t8_forest/t8_forest_partition.h b/src/t8_forest/t8_forest_partition.h
index 0f96605239..190119fe9b 100644
--- a/src/t8_forest/t8_forest_partition.h
+++ b/src/t8_forest/t8_forest_partition.h
@@ -32,12 +32,6 @@
 #include <t8.h>
 #include <t8_forest/t8_forest_general.h>
 
-/**
- * The prototype a weight function for the t8_forest_partition algorithm.
- * The function should be pure, and return a positive weight given a forest, a local tree index and an element index within the local tree
- */
-using weight_fcn_t = double (t8_forest_t, t8_locidx_t, t8_locidx_t);
-
 T8_EXTERN_C_BEGIN ();
 /**
  * Populate a forest with the partitioned elements of forest->set_from.
@@ -49,7 +43,7 @@ T8_EXTERN_C_BEGIN ();
  * \note If \a weight_callback is null, then all the elements are assumed to have the same weight
 */
 void
-t8_forest_partition (t8_forest_t forest, weight_fcn_t *weight_callback = nullptr);
+t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback);
 
 /** Create the element_offset array of a partitioned forest.
  * \param [in,out]  forest The forest.
diff --git a/src/t8_forest/t8_forest_types.h b/src/t8_forest/t8_forest_types.h
index ab2f682f61..56f0816992 100644
--- a/src/t8_forest/t8_forest_types.h
+++ b/src/t8_forest/t8_forest_types.h
@@ -46,7 +46,7 @@ typedef struct t8_forest_ghost *t8_forest_ghost_t; /**< Defined below */
  * The latter 3 can be combined, in which case the order is
  * 1. Adapt, 2. Partition, 3. Balance.
  * We store the methods in an int8_t and use these defines to
- * distinguish between them. 
+ * distinguish between them.
  */
 typedef int8_t t8_forest_from_t;
 
@@ -72,6 +72,7 @@ typedef struct t8_forest
   int set_level;          /**< Level to use in new construction. */
   int set_for_coarsening; /**< Change partition to allow
                                                      for one round of coarsening */
+  t8_weight_fcn_t *weight_function; /**< Pointer to user defined element weight function. Can be null. */
 
   sc_MPI_Comm mpicomm; /**< MPI communicator to use. */
   t8_cmesh_t cmesh;    /**< Coarse mesh to use. */
@@ -83,7 +84,7 @@ typedef struct t8_forest
   int dimension;             /**< Dimension inferred from \b cmesh. */
   int incomplete_trees;      /**< Flag to check whether the forest has (potential) incomplete trees.
                                              A tree is incomplete if an element has been removed from it.
-                                             Once an element got removed, the flag sets to 1 (true) and stays. 
+                                             Once an element got removed, the flag sets to 1 (true) and stays.
                                              For a committed forest this flag is either true on all ranks or
                                              false on all ranks. */
 
@@ -108,8 +109,8 @@ typedef struct t8_forest
   int mpisize;                    /**< Number of MPI processes. */
   int mpirank;                    /**< Number of this MPI process. */
 
-  t8_gloidx_t first_local_tree;       /**< The global index of the first local tree on this process. 
-                                             If first_local_tree is larger than last_local_tree then 
+  t8_gloidx_t first_local_tree;       /**< The global index of the first local tree on this process.
+                                             If first_local_tree is larger than last_local_tree then
                                              this processor/forest is empty.
                                              See https://github.com/DLR-AMR/t8code/wiki/Tree-indexing */
   t8_gloidx_t last_local_tree;        /**< The global index of the last local tree on this process.
@@ -190,9 +191,9 @@ typedef struct t8_profile
 
 } t8_profile_struct_t;
 
-/** 
+/**
  * This struct stores various information about a forest's ghost elements and ghost trees.
- * 
+ *
  */
 typedef struct t8_forest_ghost
 {
diff --git a/test/t8_forest/t8_gtest_forest_commit.cxx b/test/t8_forest/t8_gtest_forest_commit.cxx
index ba9ead557b..dec55971dc 100644
--- a/test/t8_forest/t8_gtest_forest_commit.cxx
+++ b/test/t8_forest/t8_gtest_forest_commit.cxx
@@ -104,7 +104,7 @@ t8_test_forest_commit_abp (t8_forest_t forest, int maxlevel)
 #if T8_TEST_LEVEL_INT < 2
   t8_forest_set_balance (forest_ada_bal_par, NULL, 0);
 #endif
-  t8_forest_set_partition (forest_ada_bal_par, NULL, 0);
+  t8_forest_set_partition (forest_ada_bal_par, NULL, 0, nullptr);
   t8_forest_commit (forest_ada_bal_par);
 
   return forest_ada_bal_par;
@@ -137,9 +137,9 @@ t8_test_forest_commit_abp_3step (t8_forest_t forest, int maxlevel)
 
   /* partition the forest */
 #if T8_TEST_LEVEL_INT < 2
-  t8_forest_set_partition (forest_partition, forest_balance, 0);
+  t8_forest_set_partition (forest_partition, forest_balance, 0, nullptr);
 #else
-  t8_forest_set_partition (forest_partition, forest_adapt, 0);
+  t8_forest_set_partition (forest_partition, forest_adapt, 0, nullptr);
 #endif
   t8_forest_commit (forest_partition);
 
diff --git a/test/t8_forest/t8_gtest_partition_data.cxx b/test/t8_forest/t8_gtest_partition_data.cxx
index 53005d74f7..2b8d2cfe6d 100644
--- a/test/t8_forest/t8_gtest_partition_data.cxx
+++ b/test/t8_forest/t8_gtest_partition_data.cxx
@@ -254,7 +254,7 @@ TEST_P (t8_test_partition_data_test, test_partition_data)
   t8_forest_t partitioned_forest;
   t8_forest_init (&partitioned_forest);
   const int partition_for_coarsening = 0;
-  t8_forest_set_partition (partitioned_forest, initial_forest, partition_for_coarsening);
+  t8_forest_set_partition (partitioned_forest, initial_forest, partition_for_coarsening, nullptr);
   t8_forest_commit (partitioned_forest);
 
   /* Test the exemplary partition_data with some arithmetic data types as well as with a custom struct. */
diff --git a/test/t8_forest_incomplete/t8_gtest_empty_global_tree.cxx b/test/t8_forest_incomplete/t8_gtest_empty_global_tree.cxx
index dc8489bc16..9db23eca1b 100644
--- a/test/t8_forest_incomplete/t8_gtest_empty_global_tree.cxx
+++ b/test/t8_forest_incomplete/t8_gtest_empty_global_tree.cxx
@@ -28,11 +28,11 @@
 #include <t8_schemes/t8_default/t8_default.hxx>
 #include <test/t8_gtest_macros.hxx>
 
-/** In this test, we are given a forest with 3 global trees. 
- * We adapt the forest so that all 6 compositions of empty 
- * global trees are the result of it. 
+/** In this test, we are given a forest with 3 global trees.
+ * We adapt the forest so that all 6 compositions of empty
+ * global trees are the result of it.
  * Therefore, \a testcase runs from 0 to 5.
- * We do this twice. Once we partition the forest in the same call. 
+ * We do this twice. Once we partition the forest in the same call.
  * The second time, we do the adapting and partitioning separately.
  * The two resulting forests must be equal.
  * */
@@ -115,11 +115,11 @@ t8_adapt_forest (t8_forest_t forest_from, t8_forest_adapt_t adapt_fn, int do_ada
   if (do_adapt) {
     t8_forest_set_adapt (forest_new, forest_from, adapt_fn, 0);
     if (do_partition) {
-      t8_forest_set_partition (forest_new, NULL, 0);
+      t8_forest_set_partition (forest_new, NULL, 0, nullptr);
     }
   }
   else if (do_partition) {
-    t8_forest_set_partition (forest_new, forest_from, 0);
+    t8_forest_set_partition (forest_new, forest_from, 0, nullptr);
   }
   if (user_data != NULL) {
     t8_forest_set_user_data (forest_new, user_data);
diff --git a/test/t8_forest_incomplete/t8_gtest_empty_local_tree.cxx b/test/t8_forest_incomplete/t8_gtest_empty_local_tree.cxx
index 80316d8e1b..e5165bef08 100644
--- a/test/t8_forest_incomplete/t8_gtest_empty_local_tree.cxx
+++ b/test/t8_forest_incomplete/t8_gtest_empty_local_tree.cxx
@@ -31,11 +31,11 @@
 
 #define MAX_NUM_RANKS 8
 
-/* In this test, a partitioned forest with one global tree and at 
- * least so many elements, such that each process has at least one 
+/* In this test, a partitioned forest with one global tree and at
+ * least so many elements, such that each process has at least one
  * local element is given. Let x be the number of mpi ranks.
  * There are 2^x many ways to empty these x local trees.
- * 
+ *
  * Example:
  * x = 3
  * instances - binary representation
@@ -43,13 +43,13 @@
  *      1    -  0 1
  *      2    -  1 0
  *      3    -  1 1
- * We remove all elements from rank with id i if the i`th bit 
+ * We remove all elements from rank with id i if the i`th bit
  * in the current instances is 0.
- * 
+ *
  * Note, this test runs only on two to maxmal 8 ranks.
- * 
- * We adapt the given forest twice. 
- * The first time, we partition the forest in the same call. 
+ *
+ * We adapt the given forest twice.
+ * The first time, we partition the forest in the same call.
  * The second time, we do the adapting and partitioning separately.
  * The two resulting forests must be equal.
  */
@@ -82,7 +82,7 @@ class DISABLED_local_tree: public testing::TestWithParam<t8_eclass_t> {
   t8_forest_t forest;
 };
 
-/** This structure contains a bitset with all 
+/** This structure contains a bitset with all
  * local trees on all processes to be removed.
  */
 struct t8_trees_to_remove
@@ -90,7 +90,7 @@ struct t8_trees_to_remove
   std::bitset<MAX_NUM_RANKS> remove;
 };
 
-/** Remove every element of rank i if the i`th bit in 
+/** Remove every element of rank i if the i`th bit in
  * the current instance \a remove is 0. */
 static int
 t8_adapt_remove (t8_forest_t forest, t8_forest_t forest_from, [[maybe_unused]] t8_locidx_t which_tree,
@@ -114,11 +114,11 @@ t8_adapt_forest (t8_forest_t forest_from, t8_forest_adapt_t adapt_fn, int do_ada
   if (do_adapt) {
     t8_forest_set_adapt (forest_new, forest_from, adapt_fn, 0);
     if (do_partition) {
-      t8_forest_set_partition (forest_new, NULL, 0);
+      t8_forest_set_partition (forest_new, NULL, 0, nullptr);
     }
   }
   else if (do_partition) {
-    t8_forest_set_partition (forest_new, forest_from, 0);
+    t8_forest_set_partition (forest_new, forest_from, 0, nullptr);
   }
   if (user_data != NULL) {
     t8_forest_set_user_data (forest_new, user_data);
diff --git a/test/t8_forest_incomplete/t8_gtest_iterate_replace.cxx b/test/t8_forest_incomplete/t8_gtest_iterate_replace.cxx
index 2f9ce5e483..f97d59273f 100644
--- a/test/t8_forest_incomplete/t8_gtest_iterate_replace.cxx
+++ b/test/t8_forest_incomplete/t8_gtest_iterate_replace.cxx
@@ -71,7 +71,7 @@ struct t8_return_data
   int *callbacks;
 };
 
-/** Inside the callback of iterate_replace we compare \a refine 
+/** Inside the callback of iterate_replace we compare \a refine
  * with the according return value of the callback of forest_adapt.
  * If true, we check the parameter \a num_outgoing, \a first_outgoing
  * \a num_incoming and \a first_incoming for correctness. */
@@ -210,11 +210,11 @@ t8_adapt_forest (t8_forest_t forest_from, t8_forest_adapt_t adapt_fn, int do_ada
   if (do_adapt) {
     t8_forest_set_adapt (forest_new, forest_from, adapt_fn, 0);
     if (do_partition) {
-      t8_forest_set_partition (forest_new, NULL, 0);
+      t8_forest_set_partition (forest_new, NULL, 0, nullptr);
     }
   }
   else if (do_partition) {
-    t8_forest_set_partition (forest_new, forest_from, 0);
+    t8_forest_set_partition (forest_new, forest_from, 0, nullptr);
   }
   if (user_data != NULL) {
     t8_forest_set_user_data (forest_new, user_data);
diff --git a/test/t8_schemes/t8_gtest_element_ref_coords.cxx b/test/t8_schemes/t8_gtest_element_ref_coords.cxx
index 6fba562446..f944fd53bb 100644
--- a/test/t8_schemes/t8_gtest_element_ref_coords.cxx
+++ b/test/t8_schemes/t8_gtest_element_ref_coords.cxx
@@ -234,7 +234,7 @@ class class_ref_coords: public testing::TestWithParam<std::tuple<t8_eclass_t, in
     t8_cmesh_t cmesh = t8_cmesh_new_from_class (eclass, sc_MPI_COMM_WORLD);
     forest = t8_forest_new_uniform (cmesh, t8_scheme_new_default (), level, 0, sc_MPI_COMM_WORLD);
     t8_forest_init (&forest_partition);
-    t8_forest_set_partition (forest_partition, forest, 0);
+    t8_forest_set_partition (forest_partition, forest, 0, nullptr);
     t8_forest_commit (forest_partition);
     forest = forest_partition;
   }
diff --git a/tutorials/features/t8_features_curved_meshes.cxx b/tutorials/features/t8_features_curved_meshes.cxx
index 67f5614a2d..d425141110 100644
--- a/tutorials/features/t8_features_curved_meshes.cxx
+++ b/tutorials/features/t8_features_curved_meshes.cxx
@@ -23,7 +23,7 @@
 /* See also: https://github.com/DLR-AMR/t8code/wiki/Feature---Curved-meshes
  *
  * This is a feature tutorial about curved meshes.
- * In the following we will generate an input mesh and input geometry in Gmsh. 
+ * In the following we will generate an input mesh and input geometry in Gmsh.
  * Furthermore, we will read those files in and build a curved forest with them.
  * As refinement criterion we will define a wall which is moving through the mesh
  * and we will refine elements based on their neighboring geometrical geometries.
@@ -51,7 +51,7 @@
 #include <string>                 /* std::string */
 #include <array>                  /* std::array */
 
-/* We use this data to control to which level the elements at which 
+/* We use this data to control to which level the elements at which
  * geometry get refined. */
 struct t8_naca_geometry_adapt_data
 {
@@ -61,7 +61,7 @@ struct t8_naca_geometry_adapt_data
   int *levels;      /** Array with refinement levels */
 };
 
-/** 
+/**
  * The adaptation callback function. This function will be called once for each element
  * and the return value decides whether this element should be refined or not.
  *   return > 0 -> This element should get refined.
@@ -72,11 +72,11 @@ struct t8_naca_geometry_adapt_data
  *   return > 0 -> The first element should get refined.
  *   return = 0 -> The first element should not get refined.
  *   return < 0 -> The whole family should get coarsened.
- * 
+ *
  * In this case, the function retrieves the geometry information of the tree the element belongs to.
  * Based on that the function looks whether the tree is linked to a specific geometry
  * and if this element touches this geometry. If true, it returns 1. Otherwise it returns 0.
- *  
+ *
  * \param [in] forest       The current forest that is in construction.
  * \param [in] forest_from  The forest from which we adapt the current forest (in our case, the uniform forest)
  * \param [in] which_tree   The process local id of the current tree.
@@ -130,9 +130,9 @@ t8_naca_geometry_adapt_callback (t8_forest_t forest, t8_forest_t forest_from, t8
   return 0;
 }
 
-/** 
+/**
  * The geometry refinement function. Here, we refine all elements, which touch certain geometries.
- *  
+ *
  * \param [in] forest           The forest that has to be refined
  * \param [in] fileprefix       The prefix of the msh and brep file.
  *                              Influences only the vtu file name.
@@ -172,7 +172,7 @@ t8_naca_geometry_refinement (t8_forest_t forest, const std::string &fileprefix,
     geometries.data (), /* Array with geometry indices */
     levels.data ()      /* Array with refinement levels */
   };
-  /* Adapt and balance the forest. 
+  /* Adapt and balance the forest.
    * Note, that we have to hand the adapt data to the forest before the commit. */
   t8_forest_init (&forest_new);
   t8_forest_set_adapt (forest_new, forest, t8_naca_geometry_adapt_callback, 1);
@@ -204,7 +204,7 @@ struct t8_naca_plane_adapt_data
   int rlevel;   /* The max refinement level */
 };
 
-/** 
+/**
  * The adaptation callback function. This function will be called once for each element
  * and the return value decides whether this element should be refined or not.
  *   return > 0 -> This element should get refined.
@@ -215,12 +215,12 @@ struct t8_naca_plane_adapt_data
  *   return > 0 -> The first element should get refined.
  *   return = 0 -> The first element should not get refined.
  *   return < 0 -> The whole family should get coarsened.
- * 
+ *
  * In this case the function checks whether the element or family is in a certain proximity to a refinement plane.
  * If true and the element does not have a max level, 1 is returned. If a family of elements
- * is too far away and the level is not below or equal the min level threshold -1 is returned. 
+ * is too far away and the level is not below or equal the min level threshold -1 is returned.
  * Otherwise 0 is returned.
- *  
+ *
  * \param [in] forest       The current forest that is in construction.
  * \param [in] forest_from  The forest from which we adapt the current forest (in our case, the uniform forest)
  * \param [in] which_tree   The process local id of the current tree.
@@ -264,10 +264,10 @@ t8_naca_plane_adapt_callback (t8_forest_t forest, t8_forest_t forest_from, t8_lo
   return 0;
 }
 
-/** 
+/**
  * The plane refinement function. Here we create a refinement loop, which moves a plane
  * through the mesh and refines it near the plane.
- *  
+ *
  * \param [in] forest           The forest which has to be refined.
  * \param [in] fileprefix       The prefix of the msh and brep file.
  *                              Influences only the vtu file name.
@@ -300,12 +300,12 @@ t8_naca_plane_refinement (t8_forest_t forest, const std::string &fileprefix, int
 
   /* Moving plane loop */
   while (adapt_data.t < steps) {
-    /* Adapt and balance the forest. 
+    /* Adapt and balance the forest.
      * Note, that we have to hand the adapt data to the forest before the commit. */
     t8_forest_init (&forest_new);
     t8_forest_set_adapt (forest_new, forest, t8_naca_plane_adapt_callback, 1);
     t8_forest_set_user_data (forest_new, &adapt_data);
-    t8_forest_set_partition (forest_new, forest, 0);
+    t8_forest_set_partition (forest_new, forest, 0, nullptr);
     t8_forest_set_balance (forest_new, forest, 0);
     t8_forest_commit (forest_new);
 
diff --git a/tutorials/general/t8_step4_partition_balance_ghost.cxx b/tutorials/general/t8_step4_partition_balance_ghost.cxx
index fffead74ba..ab37698a25 100644
--- a/tutorials/general/t8_step4_partition_balance_ghost.cxx
+++ b/tutorials/general/t8_step4_partition_balance_ghost.cxx
@@ -24,30 +24,30 @@
  *
  * This is step4 of the t8code tutorials.
  * After generating a coarse mesh (step1), building a uniform forest
- * on it (step2) and adapting this forest (step3) 
+ * on it (step2) and adapting this forest (step3)
  * we will now learn how to control the forest creation in more detail,
  * how to partition and balance a forest and how to generate a layer of ghost elements.
- * 
+ *
  * Partition: Each forest is distributed among the MPI processes. Partitioning a forest means
- *            to change this distribution in order to maintain a load balance. After we 
+ *            to change this distribution in order to maintain a load balance. After we
  *            applied partition to a forest the new forest will have equal numbers of elements
  *            on each process (+- 1).
  *            In our example we start with a uniform forest from a cmesh. This forest is constructed
- *            such that each process has the same number of elements. 
+ *            such that each process has the same number of elements.
  *            We then adapt this forest, thus refining and coarsening its elements. This changes the
  *            number of elements on each process and will almost always result in a load imbalance.
- *            You can verify this yourself by printing the process local number on each process for the 
+ *            You can verify this yourself by printing the process local number on each process for the
  *            adapted forest (for example with t8_productionf).
  *            In order to reestablish a load balance, we will construct a new forest from the adapted
  *            one via the partition feature.
- * 
- * Ghost:     Many applications require a layer of ghost elements. Ghost element of a process are 
+ *
+ * Ghost:     Many applications require a layer of ghost elements. Ghost element of a process are
  *            those elements that are not local to this process but have a (face) neighbor element that is.
  *            Telling a forest to create a ghost layer will gather all necessary information and give
  *            us access to the ghost elements.
  *            t8code does not create a layer of ghost elements by default, thus our initial uniform forest
  *            and the adapted forest will not have one.
- * 
+ *
  * Balance:   A forest fulfills the balance property if (and only if) for each element its (face) neighbors
  *            have refinement level at most +1 or -1 in comparison to the element's level.
  *            The balance property is often broken after adaptation. The Balance algorithm iterates through
@@ -65,7 +65,7 @@
  *            Note that balance changes the local number of elements and thus may also change the load balance
  *            and hence require repartitioning.
  *            Balance is usually the most expensive of t8code's mesh manipulation algorithms.
- * 
+ *
  * How you can experiment here:
  *    Partition:
  *         - Test the program with different numbers of processes and compare the local
@@ -79,7 +79,7 @@
  *           on treeid to only view those elements with treeid = -1.
  *    Balance:
  *         - View the unbalanced and balanced forest vtu files and compare them.
- *           You can color the elements by level.       
+ *           You can color the elements by level.
  *         - Exercise: Apply balance to forest_adapt and not to the twice adapted forest.
  *         - Set the no_repartition flag of t8_forest_set_balance to true and observe how the
  *           partition of the resulting forest changes.
@@ -100,25 +100,25 @@ T8_EXTERN_C_BEGIN ();
    * However, t8code offers us more control over the creation of forests.
    * For example we can control whether or not a forest should have a ghost layer,
    * be balanced/partitioned from another forest, etc.
-   * 
+   *
    * Usually, there are three steps involved in creating a forest:
    * 1. Initialize the forest with t8_forest_init.
-   *      This function will prepare a forest by setting default values for its members and 
+   *      This function will prepare a forest by setting default values for its members and
    *      initializing necessary structures.
    * 2. Set all properties that the forest should have.
    *      Here we can for example set a cmesh and refinement scheme or specify that the
    *      forest should be adapted/partitioned/balanced from another forest etc.
    *      See the t8_forest_set functions in t8_forest.h
    *      The order in which you call the set functions does not matter.
-   * 3. Commit the forest with t8_forest_commit. 
+   * 3. Commit the forest with t8_forest_commit.
    *      In this step the forest is actually created after setting all
-   *      desired properties. The forest cannot be changed after it was committed.   
-   * 
+   *      desired properties. The forest cannot be changed after it was committed.
+   *
    * The t8_forest_new functions are just wrappers around this process.
    */
 
 /* In this function we create a new forest that repartitions a given forest
- * and has a layer of ghost elements. 
+ * and has a layer of ghost elements.
  */
 static t8_forest_t
 t8_step4_partition_ghost (t8_forest_t forest)
@@ -135,17 +135,17 @@ t8_step4_partition_ghost (t8_forest_t forest)
    * This will change the distribution of the forest elements among the processes
    * in such a way that afterwards each process has the same number of elements
    * (+- 1 if the number of elements is not divisible by the number of processes).
-   * 
+   *
    * The third 0 argument is the flag 'partition_for_coarsening' which is currently not
    * implemented. Once it is, this will ensure that a family of elements will not be split
    * across multiple processes and thus one level coarsening is always possible (see also the
    * comments on coarsening in t8_step3).
    */
-  t8_forest_set_partition (new_forest, forest, 0);
+  t8_forest_set_partition (new_forest, forest, 0, nullptr);
   /* Tell the new_forest to create a ghost layer.
    * This will gather those face neighbor elements of process local element that reside
    * on a different process.
-   * 
+   *
    * We currently support ghost mode T8_GHOST_FACES that creates face neighbor ghost elements
    * and will in future also support other modes for edge/vertex neighbor ghost elements.
    */
@@ -156,7 +156,7 @@ t8_step4_partition_ghost (t8_forest_t forest)
   return new_forest;
 }
 
-/* In this function we adapt a forest as in step3 and balance it. 
+/* In this function we adapt a forest as in step3 and balance it.
  * In our main program the input forest is already adapted and then the resulting twice adapted forest will be unbalanced.
  */
 static t8_forest_t
diff --git a/tutorials/general/t8_step5_element_data.cxx b/tutorials/general/t8_step5_element_data.cxx
index 63402ca25f..6a809bc149 100644
--- a/tutorials/general/t8_step5_element_data.cxx
+++ b/tutorials/general/t8_step5_element_data.cxx
@@ -23,22 +23,22 @@
 /* See also: https://github.com/DLR-AMR/t8code/wiki/Step-5---Store-element-data
  *
  * This is step5 of the t8code tutorials.
- * In the following we will store data in the individual elements of our forest. 
- * To do this, we will again create a uniform forest, which will get adapted as in step4, 
+ * In the following we will store data in the individual elements of our forest.
+ * To do this, we will again create a uniform forest, which will get adapted as in step4,
  * with the difference that we partition, balance and create ghost elements all in the same step.
- * After adapting the forest we will learn how to build a data array and gather data for 
+ * After adapting the forest we will learn how to build a data array and gather data for
  * the local elements. Furthermore, we exchange the data values of the ghost elements and
  * output the volume data to vtu.
  *
  * How you can experiment here:
  *   - Look at the paraview output files of the adapted forest.
- *     You can apply a clip filter to look into the cube. Also you can apply (in addition) 
+ *     You can apply a clip filter to look into the cube. Also you can apply (in addition)
  *     the threshold filter to display only elements with certain properties.
  *     But at first you may just want to enter the tooltip selection mode 'Hover Cells On'
  *     to display cell information when hover over them.
  *   - Change the adaptation criterion as you wish to adapt elements or families as desired.
  *   - Store even more data per element, for instance the coordinates of its midpoint.
- *     You can again apply the threshold filter to your new data. Don't forget to write the 
+ *     You can again apply the threshold filter to your new data. Don't forget to write the
  *     data into the output file.
  *  */
 
@@ -79,7 +79,7 @@ t8_step5_build_forest (sc_MPI_Comm comm, int level)
   t8_forest_init (&forest_apbg);
   t8_forest_set_user_data (forest_apbg, &adapt_data);
   t8_forest_set_adapt (forest_apbg, forest, t8_step3_adapt_callback, 0);
-  t8_forest_set_partition (forest_apbg, NULL, 0);
+  t8_forest_set_partition (forest_apbg, NULL, 0, nullptr);
   t8_forest_set_balance (forest_apbg, NULL, 0);
   t8_forest_set_ghost (forest_apbg, 1, T8_GHOST_FACES);
   t8_forest_commit (forest_apbg);
@@ -104,7 +104,7 @@ t8_step5_create_element_data (t8_forest_t forest)
 
   /* Now we need to build an array of our data that is as long as the number
    * of elements plus the number of ghosts. You can use any allocator such as
-   * new, malloc or the t8code provide allocation macro T8_ALLOC. 
+   * new, malloc or the t8code provide allocation macro T8_ALLOC.
    * Note that in the latter case you need
    * to use T8_FREE in order to free the memory.
    */
@@ -185,7 +185,7 @@ t8_step5_exchange_ghost_data (t8_forest_t forest, struct t8_step5_data_per_eleme
 }
 
 /* Write the forest as vtu and also write the element's volumes in the file.
- * 
+ *
  * t8code supports writing element based data to vtu as long as its stored
  * as doubles. Each of the data fields to write has to be provided in its own
  * array of length num_local_elements.
@@ -215,7 +215,7 @@ t8_step5_output_data_to_vtu (t8_forest_t forest, struct t8_step5_data_per_elemen
   }
   {
     /* To write user defined data, we need the extended output function t8_forest_vtk_write_file
-     * from t8_forest_vtk.h. Despite writing user data, it also offers more control over which 
+     * from t8_forest_vtk.h. Despite writing user data, it also offers more control over which
      * properties of the forest to write. */
     int write_treeid = 1;
     int write_mpirank = 1;
diff --git a/tutorials/general/t8_step5_element_data_c_interface.c b/tutorials/general/t8_step5_element_data_c_interface.c
index ffd22b3e75..7531edc560 100644
--- a/tutorials/general/t8_step5_element_data_c_interface.c
+++ b/tutorials/general/t8_step5_element_data_c_interface.c
@@ -23,22 +23,22 @@
 /* See also: https://github.com/DLR-AMR/t8code/wiki/Step-5---Store-element-data
  *
  * This is step5 of the t8code tutorials using the C interface of t8code.
- * In the following we will store data in the individual elements of our forest. 
- * To do this, we will again create a uniform forest, which will get adapted as in step4, 
+ * In the following we will store data in the individual elements of our forest.
+ * To do this, we will again create a uniform forest, which will get adapted as in step4,
  * with the difference that we partition, balance and create ghost elements all in the same step.
- * After adapting the forest we will learn how to build a data array and gather data for 
+ * After adapting the forest we will learn how to build a data array and gather data for
  * the local elements. Furthermore, we exchange the data values of the ghost elements and
  * output the volume data to vtu.
  *
  * How you can experiment here:
  *   - Look at the paraview output files of the adapted forest.
- *     You can apply a clip filter to look into the cube. Also you can apply (in addition) 
+ *     You can apply a clip filter to look into the cube. Also you can apply (in addition)
  *     the threshold filter to display only elements with certain properties.
  *     But at first you may just want to enter the tooltip selection mode 'Hover Cells On'
  *     to display cell information when hover over them.
  *   - Change the adaptation criterion as you wish to adapt elements or families as desired.
  *   - Store even more data per element, for instance the coordinates of its midpoint.
- *     You can again apply the threshold filter to your new data. Don't forget to write the 
+ *     You can again apply the threshold filter to your new data. Don't forget to write the
  *     data into the output file.
  *  */
 
@@ -81,7 +81,7 @@ t8_step5_build_forest (sc_MPI_Comm comm, int level)
   t8_forest_init (&forest_apbg);
   t8_forest_set_user_data (forest_apbg, &adapt_data);
   t8_forest_set_adapt (forest_apbg, forest, t8_step3_adapt_callback, 0);
-  t8_forest_set_partition (forest_apbg, NULL, 0);
+  t8_forest_set_partition (forest_apbg, NULL, 0, nullptr);
   t8_forest_set_balance (forest_apbg, NULL, 0);
   t8_forest_set_ghost (forest_apbg, 1, T8_GHOST_FACES);
   t8_forest_commit (forest_apbg);
@@ -106,7 +106,7 @@ t8_step5_create_element_data (t8_forest_t forest)
 
   /* Now we need to build an array of our data that is as long as the number
    * of elements plus the number of ghosts. You can use any allocator such as
-   * new, malloc or the t8code provide allocation macro T8_ALLOC. 
+   * new, malloc or the t8code provide allocation macro T8_ALLOC.
    * Note that in the latter case you need
    * to use T8_FREE in order to free the memory.
    */
@@ -188,7 +188,7 @@ t8_step5_exchange_ghost_data (t8_forest_t forest, struct t8_step5_data_per_eleme
 }
 
 /* Write the forest as vtu and also write the element's volumes in the file.
- * 
+ *
  * t8code supports writing element based data to vtu as long as its stored
  * as doubles. Each of the data fields to write has to be provided in its own
  * array of length num_local_elements.
@@ -218,7 +218,7 @@ t8_step5_output_data_to_vtu (t8_forest_t forest, struct t8_step5_data_per_elemen
   }
   {
     /* To write user defined data, we need to extended output function t8_forest_vtk_write_file
-     * from t8_forest_vtk.h. Despite writin user data, it also offers more control over which 
+     * from t8_forest_vtk.h. Despite writin user data, it also offers more control over which
      * properties of the forest to write. */
     int write_treeid = 1;
     int write_mpirank = 1;
diff --git a/tutorials/general/t8_step6_stencil.cxx b/tutorials/general/t8_step6_stencil.cxx
index 2134ae45c5..0e191c17f8 100644
--- a/tutorials/general/t8_step6_stencil.cxx
+++ b/tutorials/general/t8_step6_stencil.cxx
@@ -23,10 +23,10 @@
 /* See also: https://github.com/DLR-AMR/t8code/wiki/Step-6-Computing-stencils
  *
  * This is step6 of the t8code tutorials using the C++ interface of t8code.
- * In the following we will store data in the individual elements of our forest. 
- * To do this, we will create a uniform forest in 2D, which will get adapted, 
+ * In the following we will store data in the individual elements of our forest.
+ * To do this, we will create a uniform forest in 2D, which will get adapted,
  * partitioned, balanced and create ghost elements all in one go.
- * After adapting the forest we build a data array and gather data for 
+ * After adapting the forest we build a data array and gather data for
  * the local elements. Next, we exchange the data values of the ghost elements and compute
  * various stencils resp. finite differences. Finally, vtu files are stored with three
  * custom data fields.
@@ -98,7 +98,7 @@ t8_step6_build_forest (sc_MPI_Comm comm, int dim, int level)
   t8_forest_init (&forest_apbg);
   t8_forest_set_user_data (forest_apbg, &adapt_data);
   t8_forest_set_adapt (forest_apbg, forest, t8_step3_adapt_callback, 0);
-  t8_forest_set_partition (forest_apbg, NULL, 0);
+  t8_forest_set_partition (forest_apbg, NULL, 0, nullptr);
   t8_forest_set_balance (forest_apbg, NULL, 0);
   t8_forest_set_ghost (forest_apbg, 1, T8_GHOST_FACES);
   t8_forest_commit (forest_apbg);
@@ -299,7 +299,7 @@ t8_step6_exchange_ghost_data (t8_forest_t forest, struct data_per_element *data)
 }
 
 /* Write the forest as vtu and also write the element's volumes in the file.
- * 
+ *
  * t8code supports writing element based data to vtu as long as its stored
  * as doubles. Each of the data fields to write has to be provided in its own
  * array of length num_local_elements.

From 7773c4cc278f5c147eb5dafe41c881d99b3aaed2 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 28 Oct 2025 11:09:49 +0100
Subject: [PATCH 14/35] indent

---
 src/t8_forest/t8_forest.cxx       | 6 ++++--
 src/t8_forest/t8_forest_general.h | 3 ++-
 src/t8_forest/t8_forest_types.h   | 4 ++--
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/t8_forest/t8_forest.cxx b/src/t8_forest/t8_forest.cxx
index 7fe3d920b1..355630d816 100644
--- a/src/t8_forest/t8_forest.cxx
+++ b/src/t8_forest/t8_forest.cxx
@@ -2814,7 +2814,8 @@ t8_forest_set_copy (t8_forest_t forest, const t8_forest_t set_from)
 }
 
 void
-t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening, t8_weight_fcn_t* weight_callback)
+t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening,
+                         t8_weight_fcn_t *weight_callback)
 {
   T8_ASSERT (forest != NULL);
   T8_ASSERT (forest->rc.refcount > 0);
@@ -3165,7 +3166,8 @@ t8_forest_commit (t8_forest_t forest)
           /* forest_partition should not change ownership of forest->set_from */
           t8_forest_ref (forest->set_from);
         }
-        t8_forest_set_partition (forest_partition, forest->set_from, forest->set_for_coarsening, forest->weight_function);
+        t8_forest_set_partition (forest_partition, forest->set_from, forest->set_for_coarsening,
+                                 forest->weight_function);
         /* activate profiling, if this forest has profiling */
         t8_forest_set_profiling (forest_partition, forest->profile != NULL);
         /* Commit the partitioned forest */
diff --git a/src/t8_forest/t8_forest_general.h b/src/t8_forest/t8_forest_general.h
index 49bae28651..de8b15e2c8 100644
--- a/src/t8_forest/t8_forest_general.h
+++ b/src/t8_forest/t8_forest_general.h
@@ -330,7 +330,8 @@ t8_forest_get_user_function (const t8_forest_t forest);
  * this setting.
  */
 void
-t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening, t8_weight_fcn_t *weight_callback);
+t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening,
+                         t8_weight_fcn_t *weight_callback);
 
 /** Set a source forest to be balanced during commit.
  * A forest is said to be balanced if each element has face neighbors of level
diff --git a/src/t8_forest/t8_forest_types.h b/src/t8_forest/t8_forest_types.h
index 56f0816992..f94f938305 100644
--- a/src/t8_forest/t8_forest_types.h
+++ b/src/t8_forest/t8_forest_types.h
@@ -69,8 +69,8 @@ typedef struct t8_forest
 {
   t8_refcount_t rc; /**< Reference counter. */
 
-  int set_level;          /**< Level to use in new construction. */
-  int set_for_coarsening; /**< Change partition to allow
+  int set_level;                    /**< Level to use in new construction. */
+  int set_for_coarsening;           /**< Change partition to allow
                                                      for one round of coarsening */
   t8_weight_fcn_t *weight_function; /**< Pointer to user defined element weight function. Can be null. */
 

From 032622c97e44cfcd36650dbe4ed0f3f9fae46b9f Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 28 Oct 2025 11:52:40 +0100
Subject: [PATCH 15/35] bugfix: replace forest with forest_from

---
 src/t8_forest/t8_forest_partition.cxx | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index fa71ef6d5b..3b369f1f2e 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -499,9 +499,9 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
         t8_gloidx_t global_elm_idx = partition_offset;
         int i = std::ceil (partition_weight_offset * mpisize / forest_weight);
 
-        for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
-          for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
-            accumulated_weight += weight_fcn (forest, ltreeid, ielm);
+        for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
+          for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest_from, ltreeid); ++ielm) {
+            accumulated_weight += weight_fcn (forest_from, ltreeid, ielm);
             ++global_elm_idx;
             if (accumulated_weight > forest_weight * i / mpisize) {
               element_offsets[i] = global_elm_idx;

From 6eebd25f4c440ac40abcfdbf08157ca7bd43382a Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 28 Oct 2025 15:28:07 +0100
Subject: [PATCH 16/35] include weighted path in tests

---
 src/t8_forest/t8_forest_partition.cxx     |  2 +-
 test/t8_forest/t8_gtest_forest_commit.cxx | 22 +++++++++-------------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 3b369f1f2e..96f80ecfb9 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -502,11 +502,11 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
         for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
           for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest_from, ltreeid); ++ielm) {
             accumulated_weight += weight_fcn (forest_from, ltreeid, ielm);
-            ++global_elm_idx;
             if (accumulated_weight > forest_weight * i / mpisize) {
               element_offsets[i] = global_elm_idx;
               ++i;
             }
+            ++global_elm_idx;
           }
         }
       }
diff --git a/test/t8_forest/t8_gtest_forest_commit.cxx b/test/t8_forest/t8_gtest_forest_commit.cxx
index dec55971dc..aba9f5a3ff 100644
--- a/test/t8_forest/t8_gtest_forest_commit.cxx
+++ b/test/t8_forest/t8_gtest_forest_commit.cxx
@@ -114,32 +114,29 @@ t8_test_forest_commit_abp (t8_forest_t forest, int maxlevel)
 static t8_forest_t
 t8_test_forest_commit_abp_3step (t8_forest_t forest, int maxlevel)
 {
+  /* adapt the forest */
   t8_forest_t forest_adapt;
-  t8_forest_t forest_partition;
-#if T8_TEST_LEVEL_INT < 2
-  t8_forest_t forest_balance;
-  t8_forest_init (&forest_balance);
-#endif
-
   t8_forest_init (&forest_adapt);
-  t8_forest_init (&forest_partition);
-
-  /* adapt the forest */
   t8_forest_set_user_data (forest_adapt, &maxlevel);
   t8_forest_set_adapt (forest_adapt, forest, t8_test_adapt_balance, 1);
   t8_forest_commit (forest_adapt);
 
 #if T8_TEST_LEVEL_INT < 2
   /* balance the forest */
+  t8_forest_t forest_balance;
+  t8_forest_init (&forest_balance);
   t8_forest_set_balance (forest_balance, forest_adapt, 0);
   t8_forest_commit (forest_balance);
 #endif
 
-  /* partition the forest */
+  /* partition the forest with a unit weight function */
+  t8_forest_t forest_partition;
+  t8_forest_init (&forest_partition);
+  auto weight_fcn = []( t8_forest_t, t8_locidx_t, t8_locidx_t ) -> double { return 1; };
 #if T8_TEST_LEVEL_INT < 2
-  t8_forest_set_partition (forest_partition, forest_balance, 0, nullptr);
+  t8_forest_set_partition (forest_partition, forest_balance, 0, weight_fcn);
 #else
-  t8_forest_set_partition (forest_partition, forest_adapt, 0, nullptr);
+  t8_forest_set_partition (forest_partition, forest_adapt, 0, weight_fcn);
 #endif
   t8_forest_commit (forest_partition);
 
@@ -148,7 +145,6 @@ t8_test_forest_commit_abp_3step (t8_forest_t forest, int maxlevel)
 
 TEST_P (forest_commit, test_forest_commit)
 {
-
   t8_forest_t forest;
   t8_forest_t forest_ada_bal_part;
   t8_forest_t forest_abp_3part;

From 44314789a92186a1bcea0655ec98abeff3a6eacf Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 28 Oct 2025 15:29:20 +0100
Subject: [PATCH 17/35] indent

---
 test/t8_forest/t8_gtest_forest_commit.cxx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/t8_forest/t8_gtest_forest_commit.cxx b/test/t8_forest/t8_gtest_forest_commit.cxx
index aba9f5a3ff..dec312d2f3 100644
--- a/test/t8_forest/t8_gtest_forest_commit.cxx
+++ b/test/t8_forest/t8_gtest_forest_commit.cxx
@@ -132,7 +132,7 @@ t8_test_forest_commit_abp_3step (t8_forest_t forest, int maxlevel)
   /* partition the forest with a unit weight function */
   t8_forest_t forest_partition;
   t8_forest_init (&forest_partition);
-  auto weight_fcn = []( t8_forest_t, t8_locidx_t, t8_locidx_t ) -> double { return 1; };
+  auto weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1; };
 #if T8_TEST_LEVEL_INT < 2
   t8_forest_set_partition (forest_partition, forest_balance, 0, weight_fcn);
 #else

From 9557cbe25d9620f2ceef3aefef3f0589ffe73a32 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Wed, 29 Oct 2025 18:51:43 +0100
Subject: [PATCH 18/35] refactor to use allgatherv in the weighted case

---
 src/t8_forest/t8_forest_partition.cxx | 147 +++++++++++++-------------
 1 file changed, 74 insertions(+), 73 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 96f80ecfb9..9abb36a3c2 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -430,99 +430,100 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
   }
 }
 
-// Computes the weight of the forest, the weight of the local partition, and the local offset (i.e. the weight
-// of all the partitions of lower rank combined)
-// If weight_fcn is null, all the elements are assumed to be of unit weight.
-static std::tuple<double, double, double>
-t8_forest_integrate_leaf_weights (t8_forest_t forest, t8_weight_fcn_t *weight_fcn)
-{
-  T8_ASSERT (t8_forest_is_committed (forest));
-
-  if (weight_fcn == nullptr) {
-    return { t8_forest_get_global_num_leaf_elements (forest), t8_forest_get_local_num_leaf_elements (forest),
-             t8_forest_get_first_local_leaf_element_id (forest) };
-  }
-
-  double local_partition_weight = 0.;
-  for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest); ++ltreeid) {
-    for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest, ltreeid); ++ielm) {
-      local_partition_weight += weight_fcn (forest, ltreeid, ielm);
-    }
-  }
-
-  double local_partition_weight_offset = 0.;
-  sc_MPI_Scan (&local_partition_weight, &local_partition_weight_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, forest->mpicomm);
-  local_partition_weight_offset
-    -= local_partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
-
-  double forest_weight = local_partition_weight_offset + local_partition_weight;
-  sc_MPI_Bcast (&forest_weight, 1, sc_MPI_DOUBLE, forest->mpisize - 1, forest->mpicomm);
-
-  return { forest_weight, local_partition_weight, local_partition_weight_offset };
-}
-
-/* Calculate the new element_offset for forest from
- * the element in forest->set_from using the provided weight function */
+// Compute forest->element_offsets according to the weight function, if provided
 static void
 t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *weight_fcn)
 {
   T8_ASSERT (t8_forest_is_initialized (forest));
   T8_ASSERT (forest->set_from != NULL);
+  T8_ASSERT (forest->element_offsets == NULL);
 
   t8_forest_t forest_from = forest->set_from;
-  sc_MPI_Comm comm = forest->mpicomm;
-  int mpisize = forest->mpisize;
+  sc_MPI_Comm comm = forest_from->mpicomm;
+  int const mpisize = forest_from->mpisize;
+  t8_gloidx_t const global_num_leaf_elements = forest_from->global_num_leaf_elements;
 
-  T8_ASSERT (forest->element_offsets == NULL);
-  /* Set the shmem array type to comm */
+  /* Initialize the shmem array */
   t8_shmem_init (comm);
   t8_shmem_set_type (comm, T8_SHMEM_BEST_TYPE);
-  /* Initialize the shmem array */
-  t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), forest->mpisize + 1, comm);
-
-  auto const [forest_weight, partition_weight, partition_weight_offset]
-    = t8_forest_integrate_leaf_weights (forest_from, weight_fcn);
-
-  t8_gloidx_t const partition_offset = t8_forest_get_first_local_leaf_element_id (forest_from);
+  t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), mpisize + 1, comm);
 
-  if (t8_shmem_array_start_writing (forest->element_offsets)) {
-    if (forest_from->global_num_leaf_elements > 0) {
+  if (global_num_leaf_elements == 0) {
+    if (t8_shmem_array_start_writing (forest->element_offsets)) {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
+      std::fill_n (element_offsets, mpisize + 1, t8_gloidx_t { 0 });
+    }
+    t8_shmem_array_end_writing (forest->element_offsets);
+    return;
+  }
 
-      if (weight_fcn == nullptr) {  // in this case forest_weight = global_num_leaf_elements
-        for (int i = 0; i < mpisize; i++) {
-          element_offsets[i] = static_cast<t8_gloidx_t> (forest_weight * i / mpisize);
-        }
-      }
-      else {
-        double accumulated_weight = partition_weight_offset;
-        t8_gloidx_t global_elm_idx = partition_offset;
-        int i = std::ceil (partition_weight_offset * mpisize / forest_weight);
-
-        for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
-          for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest_from, ltreeid); ++ielm) {
-            accumulated_weight += weight_fcn (forest_from, ltreeid, ielm);
-            if (accumulated_weight > forest_weight * i / mpisize) {
-              element_offsets[i] = global_elm_idx;
-              ++i;
-            }
-            ++global_elm_idx;
-          }
-        }
+  if (not weight_fcn) {
+    if (t8_shmem_array_start_writing (forest->element_offsets)) {
+      t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
+      for (int i = 0; i < mpisize; ++i) {
+        element_offsets[i] = std::floor (static_cast<double> (global_num_leaf_elements) * i / mpisize);
       }
-      element_offsets[forest->mpisize] = forest->global_num_leaf_elements;
+      element_offsets[mpisize] = global_num_leaf_elements;
+    }
+    t8_shmem_array_end_writing (forest->element_offsets);
+    return;
+  }
 
-      for (int i = 0; i < mpisize; i++) {
-        T8_ASSERT (0 <= element_offsets[i] && element_offsets[i] < forest_from->global_num_leaf_elements);
+  double const partition_weight = [&] () {  // sum of the weights on the local partition
+    double local_sum = 0.;
+    for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
+      for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest_from, ltreeid); ++ielm) {
+        local_sum += weight_fcn (forest_from, ltreeid, ielm);
       }
     }
-    else {
-      t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
-      for (int i = 0; i <= mpisize; i++) {
-        element_offsets[i] = 0;
+    return local_sum;
+  }();
+
+  double const partition_weight_offset = [&] () {  // partial sum of the partition weights, excluding the local rank
+    double local_offset = 0.;
+    sc_MPI_Scan (&partition_weight, &local_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, comm);
+    return local_offset
+           - partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
+  }();
+
+  double const forest_weight = [&] () {  // complete sum of the partition weights
+    double total_weight = partition_weight_offset + partition_weight;
+    sc_MPI_Bcast (&total_weight, 1, sc_MPI_DOUBLE, mpisize - 1, comm);
+    return total_weight;
+  }();
+
+  // The [rank_begin, rank_end) slice of the new offsets land in the local partition
+  int const rank_begin = std::ceil (mpisize * partition_weight_offset / forest_weight);
+  int const rank_end = std::ceil (mpisize * (partition_weight_offset + partition_weight) / forest_weight);
+  std::vector<t8_gloidx_t> local_offsets { rank_end - rank_begin, 0 };
+
+  double accumulated_weight = partition_weight_offset;
+  t8_gloidx_t global_elm_idx = t8_forest_get_first_local_leaf_element_id (forest_from);
+  int i = rank_begin;
+
+  for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
+    for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest_from, ltreeid); ++ielm) {
+      T8_ASSERT (0 <= global_elm_idx && global_elm_idx < global_num_leaf_elements);
+      accumulated_weight += weight_fcn (forest_from, ltreeid, ielm);
+      while (accumulated_weight
+             > forest_weight * i / mpisize) {  // there may be empty partitions, hence while and not if
+        T8_ASSERT (rank_begin <= i && i < rank_end);
+        local_offsets[i - rank_begin] = global_elm_idx;
+        ++i;
       }
+      ++global_elm_idx;
     }
   }
+  T8_ASSERT (i == rank_end);  // i.e. local_offsets has been filled properly
+
+  t8_shmem_array_allgatherv (local_offsets.data (), local_offsets.size (), T8_MPI_GLOIDX, forest->element_offsets,
+                             T8_MPI_GLOIDX, comm);
+
+  if (t8_shmem_array_start_writing (forest->element_offsets)) {
+    t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
+    element_offsets[0] = 0;
+    element_offsets[mpisize] = global_num_leaf_elements;
+  }
   t8_shmem_array_end_writing (forest->element_offsets);
 }
 

From 0b94ce99d78ca2184d705013159ea329cb3a3292 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Thu, 13 Nov 2025 13:51:49 +0100
Subject: [PATCH 19/35] Update src/t8_forest/t8_forest_partition.cxx

Co-authored-by: spenke91 <thomas.spenke@dlr.de>
---
 src/t8_forest/t8_forest_partition.cxx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 9abb36a3c2..d55e7029ff 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -495,7 +495,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
   // The [rank_begin, rank_end) slice of the new offsets land in the local partition
   int const rank_begin = std::ceil (mpisize * partition_weight_offset / forest_weight);
   int const rank_end = std::ceil (mpisize * (partition_weight_offset + partition_weight) / forest_weight);
-  std::vector<t8_gloidx_t> local_offsets { rank_end - rank_begin, 0 };
+  std::vector<t8_gloidx_t> local_offsets ( rank_end - rank_begin, 0 );
 
   double accumulated_weight = partition_weight_offset;
   t8_gloidx_t global_elm_idx = t8_forest_get_first_local_leaf_element_id (forest_from);

From 2fa7769dc8b90535a42459a962d1ff7885ee0809 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Fri, 14 Nov 2025 10:09:07 +0100
Subject: [PATCH 20/35] bugfix

---
 src/t8_forest/t8_forest_partition.cxx | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index d55e7029ff..69379f0f71 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -440,6 +440,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
 
   t8_forest_t forest_from = forest->set_from;
   sc_MPI_Comm comm = forest_from->mpicomm;
+  int const mpirank = forest_from->mpirank;
   int const mpisize = forest_from->mpisize;
   t8_gloidx_t const global_num_leaf_elements = forest_from->global_num_leaf_elements;
 
@@ -481,9 +482,9 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
 
   double const partition_weight_offset = [&] () {  // partial sum of the partition weights, excluding the local rank
     double local_offset = 0.;
-    sc_MPI_Scan (&partition_weight, &local_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, comm);
-    return local_offset
-           - partition_weight;  // This is semantically equivalent to calling MPI_Exscan, without the rank 0 quirks
+    double local_partition_weight = partition_weight;  // because MPI does not like const variables
+    sc_MPI_Exscan (&local_partition_weight, &local_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, comm);
+    return mpirank > 0 ? local_offset : 0;  // because the result of MPI_Exscan is undefined on rank 0
   }();
 
   double const forest_weight = [&] () {  // complete sum of the partition weights

From 349a881f918a2bd14407d68488bc54376e0763fd Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Fri, 14 Nov 2025 10:13:37 +0100
Subject: [PATCH 21/35] indent

---
 src/t8_forest/t8_forest_partition.cxx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 69379f0f71..23cfc6544e 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -496,7 +496,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
   // The [rank_begin, rank_end) slice of the new offsets land in the local partition
   int const rank_begin = std::ceil (mpisize * partition_weight_offset / forest_weight);
   int const rank_end = std::ceil (mpisize * (partition_weight_offset + partition_weight) / forest_weight);
-  std::vector<t8_gloidx_t> local_offsets ( rank_end - rank_begin, 0 );
+  std::vector<t8_gloidx_t> local_offsets (rank_end - rank_begin, 0);
 
   double accumulated_weight = partition_weight_offset;
   t8_gloidx_t global_elm_idx = t8_forest_get_first_local_leaf_element_id (forest_from);

From 8b65012955c4cb23f5ca9f0935074550f21cc12f Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Wed, 26 Nov 2025 18:10:42 +0100
Subject: [PATCH 22/35] Add suggested comment

Co-authored-by: spenke91 <thomas.spenke@dlr.de>
---
 src/t8_forest/t8_forest_partition.cxx | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 23cfc6544e..3244892b48 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -458,6 +458,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
     return;
   }
 
+  // If no weight function is provided, compute the offsets solely based on the number of elements.
   if (not weight_fcn) {
     if (t8_shmem_array_start_writing (forest->element_offsets)) {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);

From 95287fdc69caf4da77db9d3b52f80c3eab8b5ae6 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Wed, 26 Nov 2025 18:11:23 +0100
Subject: [PATCH 23/35] Add suggested comment

Co-authored-by: spenke91 <thomas.spenke@dlr.de>
---
 src/t8_forest/t8_forest_partition.cxx | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 3244892b48..e5636f16d7 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -470,7 +470,8 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
     t8_shmem_array_end_writing (forest->element_offsets);
     return;
   }
-
+// Weighted load balancing:
+// ------------------------ 
   double const partition_weight = [&] () {  // sum of the weights on the local partition
     double local_sum = 0.;
     for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {

From 0cf3d086ed2223dd8ee0c548ad5b6e0d51f99974 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Wed, 26 Nov 2025 18:11:47 +0100
Subject: [PATCH 24/35] Add suggested comment

Co-authored-by: spenke91 <thomas.spenke@dlr.de>
---
 src/t8_forest/t8_forest_partition.cxx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index e5636f16d7..b088261398 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -482,7 +482,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
     return local_sum;
   }();
 
-  double const partition_weight_offset = [&] () {  // partial sum of the partition weights, excluding the local rank
+  double const partition_weight_offset = [&] () {  // partial sum of the partition weights of all lower-rank processes (excluding the local rank)
     double local_offset = 0.;
     double local_partition_weight = partition_weight;  // because MPI does not like const variables
     sc_MPI_Exscan (&local_partition_weight, &local_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, comm);

From 0d2bf0c86538ad3fcd833da0e12e0782299d6e63 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 30 Dec 2025 13:42:56 +0100
Subject: [PATCH 25/35] fix merge bugs

---
 src/t8_forest/t8_forest_partition.cxx                | 4 ++--
 test/t8_forest/t8_gtest_partition_for_coarsening.cxx | 6 +++---
 test/t8_forest/t8_gtest_set_partition_offset.cxx     | 4 ++--
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index e61b9f74da..a615a4e994 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -97,7 +97,7 @@ t8_forest_new_gather (const t8_forest_t forest_from, const int gather_rank)
   t8_forest_init (&forest_gather);
 
   // Set partition (includes also sanity checks)
-  t8_forest_set_partition (forest_gather, forest_from, 0);
+  t8_forest_set_partition (forest_gather, forest_from, 0, nullptr);
 
   // Determine global ID of the element that will be the first local one:
   // To gather all elements on gather_rank, the first local element is set to
@@ -589,7 +589,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
         // Compute element offsets
         for (i = 0; i < mpisize; i++) {
           /* Calculate the first element index for each process. We convert to doubles to prevent overflow */
-          new_first_element_id
+          t8_gloidx_t new_first_element_id
             = (((double) i * (long double) forest_from->global_num_leaf_elements) / (double) mpisize);
           T8_ASSERT (0 <= new_first_element_id && new_first_element_id < forest_from->global_num_leaf_elements);
           element_offsets[i] = new_first_element_id;
diff --git a/test/t8_forest/t8_gtest_partition_for_coarsening.cxx b/test/t8_forest/t8_gtest_partition_for_coarsening.cxx
index 58ef50de21..30281f26b9 100644
--- a/test/t8_forest/t8_gtest_partition_for_coarsening.cxx
+++ b/test/t8_forest/t8_gtest_partition_for_coarsening.cxx
@@ -251,7 +251,7 @@ TEST_P (t8_test_partition_for_coarsening_test, test_partition_for_coarsening)
 
   // Set partitioning mode and source.
   // Note: The third argument being nonzero activates the PFC correction.
-  t8_forest_set_partition (pfc_forest, adapted_base_forest, 1);
+  t8_forest_set_partition (pfc_forest, adapted_base_forest, 1, nullptr);
 
   // Increase reference counter of base forest to avoid destruction.
   t8_forest_ref (adapted_base_forest);
@@ -291,7 +291,7 @@ TEST_P (t8_test_partition_for_coarsening_test, test_partition_for_coarsening)
   t8_forest_init (&repartitioned_coarse_pfc_forest);
 
   // Set partition mode and source.
-  t8_forest_set_partition (repartitioned_coarse_pfc_forest, coarsened_pfc_forest, 0);
+  t8_forest_set_partition (repartitioned_coarse_pfc_forest, coarsened_pfc_forest, 0, nullptr);
 
   // Commit the forest.
   t8_forest_commit (repartitioned_coarse_pfc_forest);
@@ -313,7 +313,7 @@ TEST_P (t8_test_partition_for_coarsening_test, test_partition_for_coarsening)
   // Partition without PFC option.
   t8_forest_t repartitioned_coarse_gathered_forest;
   t8_forest_init (&repartitioned_coarse_gathered_forest);
-  t8_forest_set_partition (repartitioned_coarse_gathered_forest, coarse_gathered_forest, 0);
+  t8_forest_set_partition (repartitioned_coarse_gathered_forest, coarse_gathered_forest, 0, nullptr);
   t8_forest_commit (repartitioned_coarse_gathered_forest);
 
   // (5c.) Verify that the two forests are equal:
diff --git a/test/t8_forest/t8_gtest_set_partition_offset.cxx b/test/t8_forest/t8_gtest_set_partition_offset.cxx
index d705248971..7f3c000c56 100644
--- a/test/t8_forest/t8_gtest_set_partition_offset.cxx
+++ b/test/t8_forest/t8_gtest_set_partition_offset.cxx
@@ -207,7 +207,7 @@ TEST_P (t8_test_set_partition_offset_test, test_set_partition_offset)
   t8_forest_set_partition_offset (forest_manual_partition, my_element_offset);
 
   // Create forest with custom partitioning.
-  t8_forest_set_partition (forest_manual_partition, base_forest, 0);
+  t8_forest_set_partition (forest_manual_partition, base_forest, 0, nullptr);
   t8_forest_ref (base_forest);
   t8_forest_commit (forest_manual_partition);
 
@@ -225,7 +225,7 @@ TEST_P (t8_test_set_partition_offset_test, test_set_partition_offset)
   // Repartition the manually partitioned forest
   t8_forest_t forest_repartitioned;
   t8_forest_init (&forest_repartitioned);
-  t8_forest_set_partition (forest_repartitioned, forest_manual_partition, 0);
+  t8_forest_set_partition (forest_repartitioned, forest_manual_partition, 0, nullptr);
   t8_forest_ref (forest_manual_partition);
   t8_forest_commit (forest_repartitioned);
 

From 76eabe5e5dae9ae87d98f8fda3201b79d9407249 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 30 Dec 2025 14:24:20 +0100
Subject: [PATCH 26/35] revert t8_forest_set_partition change

---
 benchmarks/t8_time_forest_partition.cxx         |  2 +-
 benchmarks/t8_time_prism_adapt.cxx              |  2 +-
 example/advect/t8_advection.cxx                 |  2 +-
 example/forest/t8_test_face_iterate.cxx         |  2 +-
 example/forest/t8_test_ghost.cxx                |  2 +-
 .../forest/t8_test_ghost_large_level_diff.cxx   |  2 +-
 example/remove/t8_example_empty_trees.cxx       |  2 +-
 src/t8_forest/t8_forest.cxx                     | 17 +++++++++++------
 src/t8_forest/t8_forest_balance.cxx             |  2 +-
 src/t8_forest/t8_forest_general.h               | 17 ++++++++++++-----
 src/t8_forest/t8_forest_partition.cxx           |  9 +++++----
 src/t8_forest/t8_forest_partition.h             | 16 ++++++----------
 test/t8_forest/t8_gtest_forest_commit.cxx       |  7 ++++---
 test/t8_forest/t8_gtest_partition_data.cxx      |  2 +-
 .../t8_gtest_partition_for_coarsening.cxx       |  6 +++---
 .../t8_forest/t8_gtest_set_partition_offset.cxx |  4 ++--
 .../t8_gtest_empty_global_tree.cxx              |  4 ++--
 .../t8_gtest_empty_local_tree.cxx               |  4 ++--
 .../t8_gtest_iterate_replace.cxx                |  4 ++--
 test/t8_schemes/t8_gtest_element_ref_coords.cxx |  2 +-
 .../features/t8_features_curved_meshes.cxx      |  2 +-
 .../t8_step4_partition_balance_ghost.cxx        |  2 +-
 tutorials/general/t8_step5_element_data.cxx     |  2 +-
 tutorials/general/t8_step6_stencil.cxx          |  2 +-
 24 files changed, 63 insertions(+), 53 deletions(-)

diff --git a/benchmarks/t8_time_forest_partition.cxx b/benchmarks/t8_time_forest_partition.cxx
index 798b5fecc8..c0ce560d7c 100644
--- a/benchmarks/t8_time_forest_partition.cxx
+++ b/benchmarks/t8_time_forest_partition.cxx
@@ -179,7 +179,7 @@ t8_time_forest_cmesh_mshfile (t8_cmesh_t cmesh, const char *vtu_prefix, sc_MPI_C
     /* partition the adapted forest */
     t8_forest_init (&forest_partition);
     /* partition the adapted forest */
-    t8_forest_set_partition (forest_partition, forest_adapt, 0, nullptr);
+    t8_forest_set_partition (forest_partition, forest_adapt, 0);
 
     /* If desired, create ghost elements and balance */
     t8_forest_set_profiling (forest_partition, 1);
diff --git a/benchmarks/t8_time_prism_adapt.cxx b/benchmarks/t8_time_prism_adapt.cxx
index 3b32c5ec4d..8853845315 100644
--- a/benchmarks/t8_time_prism_adapt.cxx
+++ b/benchmarks/t8_time_prism_adapt.cxx
@@ -128,7 +128,7 @@ t8_time_refine (int start_level, int end_level, [[maybe_unused]] int create_fore
     }
     forest_partition = forest_adapt;
     /* partition the adapted forest */
-    t8_forest_set_partition (forest_partition, NULL, 0, nullptr);
+    t8_forest_set_partition (forest_partition, NULL, 0);
     /* enable profiling for the partitioned forest */
     t8_forest_set_profiling (forest_partition, 1);
     /* if desired do balance */
diff --git a/example/advect/t8_advection.cxx b/example/advect/t8_advection.cxx
index f7b76dd53e..0853141fff 100644
--- a/example/advect/t8_advection.cxx
+++ b/example/advect/t8_advection.cxx
@@ -783,7 +783,7 @@ t8_advect_problem_partition (t8_advect_problem_t *problem, int measure_time)
   /* Enable profiling to measure runtime */
   t8_forest_set_profiling (forest_partition, 1);
   /* Partition the forest and create ghosts */
-  t8_forest_set_partition (forest_partition, problem->forest, 0, nullptr);
+  t8_forest_set_partition (forest_partition, problem->forest, 0);
   t8_forest_set_ghost (forest_partition, 1, T8_GHOST_FACES);
   t8_forest_commit (forest_partition);
 
diff --git a/example/forest/t8_test_face_iterate.cxx b/example/forest/t8_test_face_iterate.cxx
index 27e0ce1d1c..6b987f871c 100644
--- a/example/forest/t8_test_face_iterate.cxx
+++ b/example/forest/t8_test_face_iterate.cxx
@@ -138,7 +138,7 @@ t8_test_fiterate_refine_and_partition (t8_cmesh_t cmesh, int level, sc_MPI_Comm
 
   /* partition the adapted forest */
   t8_forest_init (&forest_partition);
-  t8_forest_set_partition (forest_partition, forest_adapt, 0, nullptr);
+  t8_forest_set_partition (forest_partition, forest_adapt, 0);
   t8_forest_commit (forest_partition);
   t8_debugf ("Created ghost structure with %li ghost elements.\n", (long) t8_forest_get_num_ghosts (forest_partition));
   if (!no_vtk) {
diff --git a/example/forest/t8_test_ghost.cxx b/example/forest/t8_test_ghost.cxx
index 7730d4f895..4121ba77fb 100644
--- a/example/forest/t8_test_ghost.cxx
+++ b/example/forest/t8_test_ghost.cxx
@@ -189,7 +189,7 @@ t8_test_ghost_refine_and_partition (t8_cmesh_t cmesh, const int level, sc_MPI_Co
   }
 
   /* Set the forest for partitioning */
-  t8_forest_set_partition (forest_ghost, forest, 0, nullptr);
+  t8_forest_set_partition (forest_ghost, forest, 0);
   /* Activate ghost creation */
   t8_forest_set_ghost_ext (forest_ghost, 1, T8_GHOST_FACES, ghost_version);
   /* Activate timers */
diff --git a/example/forest/t8_test_ghost_large_level_diff.cxx b/example/forest/t8_test_ghost_large_level_diff.cxx
index e79993a6e3..d6fd2c69bd 100644
--- a/example/forest/t8_test_ghost_large_level_diff.cxx
+++ b/example/forest/t8_test_ghost_large_level_diff.cxx
@@ -179,7 +179,7 @@ t8_ghost_large_level_diff (const char *prefix, int dim, int level, int refine, i
 
   /* Partition */
   t8_forest_init (&forest_partition);
-  t8_forest_set_partition (forest_partition, forest_adapt, 0, nullptr);
+  t8_forest_set_partition (forest_partition, forest_adapt, 0);
   t8_forest_set_ghost_ext (forest_partition, 1, T8_GHOST_FACES, 3);
   t8_forest_set_profiling (forest_partition, 1);
   t8_forest_commit (forest_partition);
diff --git a/example/remove/t8_example_empty_trees.cxx b/example/remove/t8_example_empty_trees.cxx
index 1c5e487b73..476f68c581 100644
--- a/example/remove/t8_example_empty_trees.cxx
+++ b/example/remove/t8_example_empty_trees.cxx
@@ -66,7 +66,7 @@ t8_strip_of_quads (t8_gloidx_t num_trees, t8_gloidx_t empty_tree, const char **v
   t8_forest_t forest_adapt;
   t8_forest_init (&forest_adapt);
   t8_forest_set_adapt (forest_adapt, forest, t8_adapt_remove, 0);
-  t8_forest_set_partition (forest_adapt, NULL, 0, nullptr);
+  t8_forest_set_partition (forest_adapt, NULL, 0);
   t8_forest_set_user_data (forest_adapt, &empty_tree);
   t8_forest_commit (forest_adapt);
 
diff --git a/src/t8_forest/t8_forest.cxx b/src/t8_forest/t8_forest.cxx
index 8ab79102d0..39c52f58ce 100644
--- a/src/t8_forest/t8_forest.cxx
+++ b/src/t8_forest/t8_forest.cxx
@@ -2808,8 +2808,7 @@ t8_forest_set_copy (t8_forest_t forest, const t8_forest_t set_from)
 }
 
 void
-t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening,
-                         t8_weight_fcn_t *weight_callback)
+t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening)
 {
   T8_ASSERT (forest != NULL);
   T8_ASSERT (forest->rc.refcount > 0);
@@ -2819,7 +2818,7 @@ t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set
   T8_ASSERT (forest->scheme == NULL);
 
   forest->set_for_coarsening = set_for_coarsening;
-  forest->weight_function = weight_callback;
+  forest->weight_function = nullptr;
 
   if (set_from != NULL) {
     /* If set_from = NULL, we assume a previous forest_from was set */
@@ -2835,6 +2834,12 @@ t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set
   }
 }
 
+void
+t8_forest_set_partition_weights (t8_forest_t forest, t8_weight_fcn_t *weight_callback)
+{
+  forest->weight_function = weight_callback;
+}
+
 void
 t8_forest_set_balance (t8_forest_t forest, const t8_forest_t set_from, int no_repartition)
 {
@@ -3160,8 +3165,8 @@ t8_forest_commit (t8_forest_t forest)
           /* forest_partition should not change ownership of forest->set_from */
           t8_forest_ref (forest->set_from);
         }
-        t8_forest_set_partition (forest_partition, forest->set_from, forest->set_for_coarsening,
-                                 forest->weight_function);
+        t8_forest_set_partition (forest_partition, forest->set_from, forest->set_for_coarsening);
+        t8_forest_set_partition_weights (forest_partition, forest->weight_function);
         /* activate profiling, if this forest has profiling */
         t8_forest_set_profiling (forest_partition, forest->profile != NULL);
         /* Commit the partitioned forest */
@@ -3182,7 +3187,7 @@ t8_forest_commit (t8_forest_t forest)
         /* Initialize the trees array of the forest */
         forest->trees = sc_array_new (sizeof (t8_tree_struct_t));
         /* partition the forest */
-        t8_forest_partition (forest, forest->weight_function);
+        t8_forest_partition (forest);
       }
     }
     if (forest->from_method & T8_FOREST_FROM_BALANCE) {
diff --git a/src/t8_forest/t8_forest_balance.cxx b/src/t8_forest/t8_forest_balance.cxx
index abdea32c93..ae67e5dc21 100644
--- a/src/t8_forest/t8_forest_balance.cxx
+++ b/src/t8_forest/t8_forest_balance.cxx
@@ -250,7 +250,7 @@ t8_forest_balance (t8_forest_t forest, int repartition)
       t8_forest_init (&forest_partition);
       /* Update the maximum occurring level */
       forest_partition->maxlevel_existing = forest_temp->maxlevel_existing;
-      t8_forest_set_partition (forest_partition, forest_temp, 0, nullptr);
+      t8_forest_set_partition (forest_partition, forest_temp, 0);
       t8_forest_set_ghost (forest_partition, 1, T8_GHOST_FACES);
       /* If profiling is enabled, measure partition runtimes */
       if (forest->profile != NULL) {
diff --git a/src/t8_forest/t8_forest_general.h b/src/t8_forest/t8_forest_general.h
index a2e04c6845..b31119eb39 100644
--- a/src/t8_forest/t8_forest_general.h
+++ b/src/t8_forest/t8_forest_general.h
@@ -320,7 +320,6 @@ t8_forest_get_user_function (const t8_forest_t forest);
  * \param [in]      set_for_coarsening CURRENTLY DISABLED. If true, then the partitions
  *                          are choose such that coarsening an element once is a process local
  *                          operation.
- * \param [in]      weight_callback A callback function defining element weights for the partitioning. If null, all the elements are assumed to have the same weight. Must be free of side-effects.
  * \note This setting can be combined with \ref t8_forest_set_adapt and \ref
  * t8_forest_set_balance. The order in which these operations are executed is always
  * 1) Adapt 2) Partition 3) Balance.
@@ -330,8 +329,16 @@ t8_forest_get_user_function (const t8_forest_t forest);
  * this setting.
  */
 void
-t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening,
-                         t8_weight_fcn_t *weight_callback);
+t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening);
+
+/** Set a user-defined weight function to guide the partitioning.
+ * \param [in, out] forest  The forest.
+ * \param [in]      weight_callback A callback function defining element weights for the partitioning.
+ * \pre \a weight_callback must be free of side-effects, the behavior is undefined otherwise
+ * \note If \a weight_callback is null, then all the elements are assumed to have the same weight
+ */
+void
+t8_forest_set_partition_weights (t8_forest_t forest, t8_weight_fcn_t *weight_callback);
 
 /** Set a source forest to be balanced during commit.
  * A forest is said to be balanced if each element has face neighbors of level
@@ -476,8 +483,8 @@ t8_forest_tree_is_local (const t8_forest_t forest, const t8_locidx_t local_tree)
  * \param [in]      forest The forest.
  * \param [in]      gtreeid The global id of a tree.
  * \return                 The tree's local id in \a forest, if it is a local tree.
- *                         A negative number if not. Ghosts trees are not considered 
- *                         as local. 
+ *                         A negative number if not. Ghosts trees are not considered
+ *                         as local.
  * \see t8_forest_get_local_or_ghost_id for ghost trees.
  * \see https://github.com/DLR-AMR/t8code/wiki/Tree-indexing for more details about tree indexing.
  */
diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index a615a4e994..a44e03142a 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -97,7 +97,7 @@ t8_forest_new_gather (const t8_forest_t forest_from, const int gather_rank)
   t8_forest_init (&forest_gather);
 
   // Set partition (includes also sanity checks)
-  t8_forest_set_partition (forest_gather, forest_from, 0, nullptr);
+  t8_forest_set_partition (forest_gather, forest_from, 0);
 
   // Determine global ID of the element that will be the first local one:
   // To gather all elements on gather_rank, the first local element is set to
@@ -474,7 +474,7 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
 
 // Compute forest->element_offsets according to the weight function, if provided
 static void
-t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *weight_fcn)
+t8_forest_partition_compute_new_offset (t8_forest_t forest)
 {
   T8_ASSERT (t8_forest_is_initialized (forest));
   T8_ASSERT (forest->set_from != NULL);
@@ -485,6 +485,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest, t8_weight_fcn_t *wei
   int const mpirank = forest_from->mpirank;
   int const mpisize = forest_from->mpisize;
   t8_gloidx_t const global_num_leaf_elements = forest_from->global_num_leaf_elements;
+  auto const weight_fcn = forest->weight_function;
 
   /* Initialize the shmem array */
   t8_shmem_init (comm);
@@ -1339,7 +1340,7 @@ t8_forest_partition_given (t8_forest_t forest, const int send_data, const sc_arr
  * Currently the elements are distributed evenly (each element has the same weight).
  */
 void
-t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback)
+t8_forest_partition (t8_forest_t forest)
 {
   t8_forest_t forest_from;
   int create_offset_from = 0;
@@ -1369,7 +1370,7 @@ t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback)
   /* TODO: if offsets already exist on forest_from, check it for consistency */
 
   /* We now calculate the new element offsets */
-  t8_forest_partition_compute_new_offset (forest, weight_callback);
+  t8_forest_partition_compute_new_offset (forest);
   t8_forest_partition_given (forest, 0, NULL, NULL);
 
   T8_ASSERT ((size_t) t8_forest_get_num_local_trees (forest_from) == forest_from->trees->elem_count);
diff --git a/src/t8_forest/t8_forest_partition.h b/src/t8_forest/t8_forest_partition.h
index b7c8115b6e..6b0bda0e43 100644
--- a/src/t8_forest/t8_forest_partition.h
+++ b/src/t8_forest/t8_forest_partition.h
@@ -36,24 +36,20 @@ T8_EXTERN_C_BEGIN ();
 
 /**
  * Populate a forest with the partitioned elements of forest->set_from.
- * Currently the elements are distributed evenly (each element has the same weight).
  *
  * \param [in,out]  forest  The forest.
- * \param [in]      weight_callback A callback function defining element weights for the partitioning
- * \pre \a weight_callback must be free of side-effects, the behavior is undefined otherwise
- * \note If \a weight_callback is null, then all the elements are assumed to have the same weight
 */
 void
-t8_forest_partition (t8_forest_t forest, t8_weight_fcn_t *weight_callback);
+t8_forest_partition (t8_forest_t forest);
 
 /**
  * Create a new forest that gathers a given forest on one process.
- * 
+ *
  * This functionality is mostly required for comparison purposes and sanity checks within the testing framework.
- * 
+ *
  * \param[in] forest_from   the forest that should be gathered on one rank
  * \param[in] gather_rank   the rank of the process the forest will be gathered on
- * 
+ *
  * \return The gathered forest: The same as \a forest_from, but all elements are on rank \a gather_rank.
 */
 t8_forest_t
@@ -61,9 +57,9 @@ t8_forest_new_gather (const t8_forest_t forest_from, const int gather_rank);
 
 /**
  * Manually set the partition offset of the current process.
- * 
+ *
  * If set, the next partitioning of the forest will use the manually defined element offsets.
- * 
+ *
  * \param[in,out] forest                the considered forest
  * \param[in]     first_global_element  the global ID that will become the first local element
 */
diff --git a/test/t8_forest/t8_gtest_forest_commit.cxx b/test/t8_forest/t8_gtest_forest_commit.cxx
index cac9241a74..da5b1e0857 100644
--- a/test/t8_forest/t8_gtest_forest_commit.cxx
+++ b/test/t8_forest/t8_gtest_forest_commit.cxx
@@ -104,7 +104,7 @@ t8_test_forest_commit_abp (t8_forest_t forest, int maxlevel)
 #if T8_TEST_LEVEL_INT < 2
   t8_forest_set_balance (forest_ada_bal_par, NULL, 0);
 #endif
-  t8_forest_set_partition (forest_ada_bal_par, NULL, 0, nullptr);
+  t8_forest_set_partition (forest_ada_bal_par, NULL, 0);
   t8_forest_commit (forest_ada_bal_par);
 
   return forest_ada_bal_par;
@@ -134,10 +134,11 @@ t8_test_forest_commit_abp_3step (t8_forest_t forest, int maxlevel)
   t8_forest_init (&forest_partition);
   auto weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1; };
 #if T8_TEST_LEVEL_INT < 2
-  t8_forest_set_partition (forest_partition, forest_balance, 0, weight_fcn);
+  t8_forest_set_partition (forest_partition, forest_balance, 0);
 #else
-  t8_forest_set_partition (forest_partition, forest_adapt, 0, weight_fcn);
+  t8_forest_set_partition (forest_partition, forest_adapt, 0);
 #endif
+  t8_forest_set_partition_weights (forest_partition, weight_fcn);
   t8_forest_commit (forest_partition);
 
   return forest_partition;
diff --git a/test/t8_forest/t8_gtest_partition_data.cxx b/test/t8_forest/t8_gtest_partition_data.cxx
index 62366be7fc..f68f1ef6b8 100644
--- a/test/t8_forest/t8_gtest_partition_data.cxx
+++ b/test/t8_forest/t8_gtest_partition_data.cxx
@@ -254,7 +254,7 @@ TEST_P (t8_test_partition_data_test, test_partition_data)
   t8_forest_t partitioned_forest;
   t8_forest_init (&partitioned_forest);
   const int partition_for_coarsening = 0;
-  t8_forest_set_partition (partitioned_forest, initial_forest, partition_for_coarsening, nullptr);
+  t8_forest_set_partition (partitioned_forest, initial_forest, partition_for_coarsening);
   t8_forest_commit (partitioned_forest);
 
   /* Test the exemplary partition_data with some arithmetic data types as well as with a custom struct. */
diff --git a/test/t8_forest/t8_gtest_partition_for_coarsening.cxx b/test/t8_forest/t8_gtest_partition_for_coarsening.cxx
index 30281f26b9..58ef50de21 100644
--- a/test/t8_forest/t8_gtest_partition_for_coarsening.cxx
+++ b/test/t8_forest/t8_gtest_partition_for_coarsening.cxx
@@ -251,7 +251,7 @@ TEST_P (t8_test_partition_for_coarsening_test, test_partition_for_coarsening)
 
   // Set partitioning mode and source.
   // Note: The third argument being nonzero activates the PFC correction.
-  t8_forest_set_partition (pfc_forest, adapted_base_forest, 1, nullptr);
+  t8_forest_set_partition (pfc_forest, adapted_base_forest, 1);
 
   // Increase reference counter of base forest to avoid destruction.
   t8_forest_ref (adapted_base_forest);
@@ -291,7 +291,7 @@ TEST_P (t8_test_partition_for_coarsening_test, test_partition_for_coarsening)
   t8_forest_init (&repartitioned_coarse_pfc_forest);
 
   // Set partition mode and source.
-  t8_forest_set_partition (repartitioned_coarse_pfc_forest, coarsened_pfc_forest, 0, nullptr);
+  t8_forest_set_partition (repartitioned_coarse_pfc_forest, coarsened_pfc_forest, 0);
 
   // Commit the forest.
   t8_forest_commit (repartitioned_coarse_pfc_forest);
@@ -313,7 +313,7 @@ TEST_P (t8_test_partition_for_coarsening_test, test_partition_for_coarsening)
   // Partition without PFC option.
   t8_forest_t repartitioned_coarse_gathered_forest;
   t8_forest_init (&repartitioned_coarse_gathered_forest);
-  t8_forest_set_partition (repartitioned_coarse_gathered_forest, coarse_gathered_forest, 0, nullptr);
+  t8_forest_set_partition (repartitioned_coarse_gathered_forest, coarse_gathered_forest, 0);
   t8_forest_commit (repartitioned_coarse_gathered_forest);
 
   // (5c.) Verify that the two forests are equal:
diff --git a/test/t8_forest/t8_gtest_set_partition_offset.cxx b/test/t8_forest/t8_gtest_set_partition_offset.cxx
index 7f3c000c56..d705248971 100644
--- a/test/t8_forest/t8_gtest_set_partition_offset.cxx
+++ b/test/t8_forest/t8_gtest_set_partition_offset.cxx
@@ -207,7 +207,7 @@ TEST_P (t8_test_set_partition_offset_test, test_set_partition_offset)
   t8_forest_set_partition_offset (forest_manual_partition, my_element_offset);
 
   // Create forest with custom partitioning.
-  t8_forest_set_partition (forest_manual_partition, base_forest, 0, nullptr);
+  t8_forest_set_partition (forest_manual_partition, base_forest, 0);
   t8_forest_ref (base_forest);
   t8_forest_commit (forest_manual_partition);
 
@@ -225,7 +225,7 @@ TEST_P (t8_test_set_partition_offset_test, test_set_partition_offset)
   // Repartition the manually partitioned forest
   t8_forest_t forest_repartitioned;
   t8_forest_init (&forest_repartitioned);
-  t8_forest_set_partition (forest_repartitioned, forest_manual_partition, 0, nullptr);
+  t8_forest_set_partition (forest_repartitioned, forest_manual_partition, 0);
   t8_forest_ref (forest_manual_partition);
   t8_forest_commit (forest_repartitioned);
 
diff --git a/test/t8_forest_incomplete/t8_gtest_empty_global_tree.cxx b/test/t8_forest_incomplete/t8_gtest_empty_global_tree.cxx
index 9db23eca1b..f7dd31cfff 100644
--- a/test/t8_forest_incomplete/t8_gtest_empty_global_tree.cxx
+++ b/test/t8_forest_incomplete/t8_gtest_empty_global_tree.cxx
@@ -115,11 +115,11 @@ t8_adapt_forest (t8_forest_t forest_from, t8_forest_adapt_t adapt_fn, int do_ada
   if (do_adapt) {
     t8_forest_set_adapt (forest_new, forest_from, adapt_fn, 0);
     if (do_partition) {
-      t8_forest_set_partition (forest_new, NULL, 0, nullptr);
+      t8_forest_set_partition (forest_new, NULL, 0);
     }
   }
   else if (do_partition) {
-    t8_forest_set_partition (forest_new, forest_from, 0, nullptr);
+    t8_forest_set_partition (forest_new, forest_from, 0);
   }
   if (user_data != NULL) {
     t8_forest_set_user_data (forest_new, user_data);
diff --git a/test/t8_forest_incomplete/t8_gtest_empty_local_tree.cxx b/test/t8_forest_incomplete/t8_gtest_empty_local_tree.cxx
index e5165bef08..d88d41845e 100644
--- a/test/t8_forest_incomplete/t8_gtest_empty_local_tree.cxx
+++ b/test/t8_forest_incomplete/t8_gtest_empty_local_tree.cxx
@@ -114,11 +114,11 @@ t8_adapt_forest (t8_forest_t forest_from, t8_forest_adapt_t adapt_fn, int do_ada
   if (do_adapt) {
     t8_forest_set_adapt (forest_new, forest_from, adapt_fn, 0);
     if (do_partition) {
-      t8_forest_set_partition (forest_new, NULL, 0, nullptr);
+      t8_forest_set_partition (forest_new, NULL, 0);
     }
   }
   else if (do_partition) {
-    t8_forest_set_partition (forest_new, forest_from, 0, nullptr);
+    t8_forest_set_partition (forest_new, forest_from, 0);
   }
   if (user_data != NULL) {
     t8_forest_set_user_data (forest_new, user_data);
diff --git a/test/t8_forest_incomplete/t8_gtest_iterate_replace.cxx b/test/t8_forest_incomplete/t8_gtest_iterate_replace.cxx
index f97d59273f..074d485ee4 100644
--- a/test/t8_forest_incomplete/t8_gtest_iterate_replace.cxx
+++ b/test/t8_forest_incomplete/t8_gtest_iterate_replace.cxx
@@ -210,11 +210,11 @@ t8_adapt_forest (t8_forest_t forest_from, t8_forest_adapt_t adapt_fn, int do_ada
   if (do_adapt) {
     t8_forest_set_adapt (forest_new, forest_from, adapt_fn, 0);
     if (do_partition) {
-      t8_forest_set_partition (forest_new, NULL, 0, nullptr);
+      t8_forest_set_partition (forest_new, NULL, 0);
     }
   }
   else if (do_partition) {
-    t8_forest_set_partition (forest_new, forest_from, 0, nullptr);
+    t8_forest_set_partition (forest_new, forest_from, 0);
   }
   if (user_data != NULL) {
     t8_forest_set_user_data (forest_new, user_data);
diff --git a/test/t8_schemes/t8_gtest_element_ref_coords.cxx b/test/t8_schemes/t8_gtest_element_ref_coords.cxx
index f944fd53bb..6fba562446 100644
--- a/test/t8_schemes/t8_gtest_element_ref_coords.cxx
+++ b/test/t8_schemes/t8_gtest_element_ref_coords.cxx
@@ -234,7 +234,7 @@ class class_ref_coords: public testing::TestWithParam<std::tuple<t8_eclass_t, in
     t8_cmesh_t cmesh = t8_cmesh_new_from_class (eclass, sc_MPI_COMM_WORLD);
     forest = t8_forest_new_uniform (cmesh, t8_scheme_new_default (), level, 0, sc_MPI_COMM_WORLD);
     t8_forest_init (&forest_partition);
-    t8_forest_set_partition (forest_partition, forest, 0, nullptr);
+    t8_forest_set_partition (forest_partition, forest, 0);
     t8_forest_commit (forest_partition);
     forest = forest_partition;
   }
diff --git a/tutorials/features/t8_features_curved_meshes.cxx b/tutorials/features/t8_features_curved_meshes.cxx
index 20d5ee9512..58eb3a622c 100644
--- a/tutorials/features/t8_features_curved_meshes.cxx
+++ b/tutorials/features/t8_features_curved_meshes.cxx
@@ -305,7 +305,7 @@ t8_naca_plane_refinement (t8_forest_t forest, const std::string &fileprefix, int
     t8_forest_init (&forest_new);
     t8_forest_set_adapt (forest_new, forest, t8_naca_plane_adapt_callback, 1);
     t8_forest_set_user_data (forest_new, &adapt_data);
-    t8_forest_set_partition (forest_new, forest, 0, nullptr);
+    t8_forest_set_partition (forest_new, forest, 0);
     t8_forest_set_balance (forest_new, forest, 0);
     t8_forest_commit (forest_new);
 
diff --git a/tutorials/general/t8_step4_partition_balance_ghost.cxx b/tutorials/general/t8_step4_partition_balance_ghost.cxx
index 48190b0acb..5092a79294 100644
--- a/tutorials/general/t8_step4_partition_balance_ghost.cxx
+++ b/tutorials/general/t8_step4_partition_balance_ghost.cxx
@@ -141,7 +141,7 @@ t8_step4_partition_ghost (t8_forest_t forest)
    * across multiple processes and thus one level coarsening is always possible (see also the
    * comments on coarsening in t8_step3).
    */
-  t8_forest_set_partition (new_forest, forest, 0, nullptr);
+  t8_forest_set_partition (new_forest, forest, 0);
   /* Tell the new_forest to create a ghost layer.
    * This will gather those face neighbor elements of process local element that reside
    * on a different process.
diff --git a/tutorials/general/t8_step5_element_data.cxx b/tutorials/general/t8_step5_element_data.cxx
index fbef32940f..200bcfe8d0 100644
--- a/tutorials/general/t8_step5_element_data.cxx
+++ b/tutorials/general/t8_step5_element_data.cxx
@@ -79,7 +79,7 @@ t8_step5_build_forest (sc_MPI_Comm comm, int level)
   t8_forest_init (&forest_apbg);
   t8_forest_set_user_data (forest_apbg, &adapt_data);
   t8_forest_set_adapt (forest_apbg, forest, t8_step3_adapt_callback, 0);
-  t8_forest_set_partition (forest_apbg, NULL, 0, nullptr);
+  t8_forest_set_partition (forest_apbg, NULL, 0);
   t8_forest_set_balance (forest_apbg, NULL, 0);
   t8_forest_set_ghost (forest_apbg, 1, T8_GHOST_FACES);
   t8_forest_commit (forest_apbg);
diff --git a/tutorials/general/t8_step6_stencil.cxx b/tutorials/general/t8_step6_stencil.cxx
index 1e2cfb2223..9804375b5c 100644
--- a/tutorials/general/t8_step6_stencil.cxx
+++ b/tutorials/general/t8_step6_stencil.cxx
@@ -98,7 +98,7 @@ t8_step6_build_forest (sc_MPI_Comm comm, int dim, int level)
   t8_forest_init (&forest_apbg);
   t8_forest_set_user_data (forest_apbg, &adapt_data);
   t8_forest_set_adapt (forest_apbg, forest, t8_step3_adapt_callback, 0);
-  t8_forest_set_partition (forest_apbg, NULL, 0, nullptr);
+  t8_forest_set_partition (forest_apbg, NULL, 0);
   t8_forest_set_balance (forest_apbg, NULL, 0);
   t8_forest_set_ghost (forest_apbg, 1, T8_GHOST_FACES);
   t8_forest_commit (forest_apbg);

From c7ae82ec73564e75e3bf5aef4786772c4236e83e Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 30 Dec 2025 14:27:23 +0100
Subject: [PATCH 27/35] add doc note

---
 src/t8_forest/t8_forest_general.h | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/t8_forest/t8_forest_general.h b/src/t8_forest/t8_forest_general.h
index b31119eb39..93611d3b10 100644
--- a/src/t8_forest/t8_forest_general.h
+++ b/src/t8_forest/t8_forest_general.h
@@ -327,6 +327,8 @@ t8_forest_get_user_function (const t8_forest_t forest);
  * false, it is not necessary to call \ref t8_forest_set_partition additionally.
  * \note This setting may not be combined with \ref t8_forest_set_copy and overwrites
  * this setting.
+ * \note Calling this function sets the forest for unweighted partitioning. If weighted
+ * partitioning is desired, \ref t8_forest_set_partition_weights must be called after.
  */
 void
 t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening);

From 98ff2de03ef614587b5a6afc96cc07e8eadd215f Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 30 Dec 2025 14:29:51 +0100
Subject: [PATCH 28/35] indent

---
 src/t8_forest/t8_forest_partition.cxx | 17 +++++++++--------
 src/t8_forest/t8_forest_types.h       |  4 ++--
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index a44e03142a..70cbdf0d91 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -513,8 +513,8 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
     t8_shmem_array_end_writing (forest->element_offsets);
     return;
   }
-// Weighted load balancing:
-// ------------------------
+  // Weighted load balancing:
+  // ------------------------
   double const partition_weight = [&] () {  // sum of the weights on the local partition
     double local_sum = 0.;
     for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
@@ -525,12 +525,13 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
     return local_sum;
   }();
 
-  double const partition_weight_offset = [&] () {  // partial sum of the partition weights of all lower-rank processes (excluding the local rank)
-    double local_offset = 0.;
-    double local_partition_weight = partition_weight;  // because MPI does not like const variables
-    sc_MPI_Exscan (&local_partition_weight, &local_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, comm);
-    return mpirank > 0 ? local_offset : 0;  // because the result of MPI_Exscan is undefined on rank 0
-  }();
+  double const partition_weight_offset
+    = [&] () {  // partial sum of the partition weights of all lower-rank processes (excluding the local rank)
+        double local_offset = 0.;
+        double local_partition_weight = partition_weight;  // because MPI does not like const variables
+        sc_MPI_Exscan (&local_partition_weight, &local_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, comm);
+        return mpirank > 0 ? local_offset : 0;  // because the result of MPI_Exscan is undefined on rank 0
+      }();
 
   double const forest_weight = [&] () {  // complete sum of the partition weights
     double total_weight = partition_weight_offset + partition_weight;
diff --git a/src/t8_forest/t8_forest_types.h b/src/t8_forest/t8_forest_types.h
index a9abffc4f5..a855e4adcd 100644
--- a/src/t8_forest/t8_forest_types.h
+++ b/src/t8_forest/t8_forest_types.h
@@ -73,8 +73,8 @@ typedef struct t8_forest
   t8_gloidx_t
     set_first_global_element; /**< If set_partition_offset is true, the global ID of the first local element after partitioning.*/
 
-  int set_level;          /**< Level to use in new construction. */
-  int set_for_coarsening; /**< Change partition to allow
+  int set_level;                    /**< Level to use in new construction. */
+  int set_for_coarsening;           /**< Change partition to allow
                                                      for one round of coarsening */
   t8_weight_fcn_t *weight_function; /**< Pointer to user defined element weight function. Can be null. */
 

From cdda894538172bc47e235a8a73a04e4bf62c9aa7 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 30 Dec 2025 15:29:30 +0100
Subject: [PATCH 29/35] fix failing test

---
 src/t8_forest/t8_forest_partition.cxx | 97 +++++++++++++--------------
 1 file changed, 47 insertions(+), 50 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 70cbdf0d91..4602cf12d3 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -32,6 +32,21 @@
 #include <t8_cmesh/t8_cmesh_internal/t8_cmesh_offset.h>
 #include <t8_schemes/t8_scheme.hxx>
 
+namespace { // to prevent this implementation detail from leaking outside of this TU
+  template< typename Callable >
+  class [[maybe_unused]] DeferToScopeExit {
+   private:
+    Callable f_;
+   public:
+    DeferToScopeExit( Callable&& f ): f_{ f } {}
+    DeferToScopeExit( DeferToScopeExit const& ) = delete;
+    DeferToScopeExit( DeferToScopeExit&& ) noexcept = delete;
+    DeferToScopeExit& operator=( DeferToScopeExit const& ) = delete;
+    DeferToScopeExit& operator=( DeferToScopeExit&& ) noexcept = delete;
+    ~DeferToScopeExit(){ f_(); }
+  };
+}
+
 /* We want to export the whole implementation to be callable from "C" */
 T8_EXTERN_C_BEGIN ();
 
@@ -492,6 +507,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   t8_shmem_set_type (comm, T8_SHMEM_BEST_TYPE);
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), mpisize + 1, comm);
 
+  // Get rid upfront of the empty forest corner case
   if (global_num_leaf_elements == 0) {
     if (t8_shmem_array_start_writing (forest->element_offsets)) {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
@@ -501,6 +517,34 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
     return;
   }
 
+  // If the partition-for-coarsening flag is set, register the partitioning correction to be exectuted on
+  // before returning. This correction prevents families of elements being split across process boundaries.
+  auto pfc = DeferToScopeExit{
+    [&](){
+      if (forest->set_for_coarsening) {
+        t8_forest_pfc_correction_offsets (forest);
+      }
+    }
+  };
+
+  // If a custom partitioning is set, all-gather the manually set element offsets.
+  if (forest->set_partition_offset != 0) {
+    std::vector<t8_gloidx_t> custom_element_offsets(forest->mpisize, 0);
+    const int retval = sc_MPI_Allgather (&forest->set_first_global_element, 1, T8_MPI_GLOIDX,
+                                         custom_element_offsets.data (), 1, T8_MPI_GLOIDX, comm);
+    SC_CHECK_MPI (retval);
+
+    if (t8_shmem_array_start_writing (forest->element_offsets)) {
+      t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
+      // Apply manually set element offsets.
+      std::copy_n (custom_element_offsets.data (), static_cast<size_t> (mpisize), element_offsets);
+      // The last entry is the same in both cases.
+      element_offsets[forest->mpisize] = forest->global_num_leaf_elements;
+    }
+    t8_shmem_array_end_writing (forest->element_offsets);
+    return;
+  }
+
   // If no weight function is provided, compute the offsets solely based on the number of elements.
   if (not weight_fcn) {
     if (t8_shmem_array_start_writing (forest->element_offsets)) {
@@ -566,59 +610,12 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   t8_shmem_array_allgatherv (local_offsets.data (), local_offsets.size (), T8_MPI_GLOIDX, forest->element_offsets,
                              T8_MPI_GLOIDX, comm);
 
-  // Define vector of manually set element offsets.
-  std::vector<t8_gloidx_t> custom_element_offsets;
-
-  // If a custom partitioning is set, all-gather the manually set element offsets.
-  if (forest->set_partition_offset != 0) {
-    custom_element_offsets.resize (forest->mpisize);
-    const int retval = sc_MPI_Allgather (&forest->set_first_global_element, 1, T8_MPI_GLOIDX,
-                                         custom_element_offsets.data (), 1, T8_MPI_GLOIDX, comm);
-    SC_CHECK_MPI (retval);
-  }
-
   if (t8_shmem_array_start_writing (forest->element_offsets)) {
-    // t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
-    // element_offsets[0] = 0;
-    // element_offsets[mpisize] = global_num_leaf_elements;
-    if (forest_from->global_num_leaf_elements > 0) {
-
-      t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
-
-      // Distinguish whether custom element offsets were set:
-      if (forest->set_partition_offset == 0) {
-
-        // Compute element offsets
-        for (i = 0; i < mpisize; i++) {
-          /* Calculate the first element index for each process. We convert to doubles to prevent overflow */
-          t8_gloidx_t new_first_element_id
-            = (((double) i * (long double) forest_from->global_num_leaf_elements) / (double) mpisize);
-          T8_ASSERT (0 <= new_first_element_id && new_first_element_id < forest_from->global_num_leaf_elements);
-          element_offsets[i] = new_first_element_id;
-        }
-      }
-      else {
-
-        // Apply manually set element offsets.
-        std::copy_n (custom_element_offsets.data (), static_cast<size_t> (mpisize), element_offsets);
-      }
-      // The last entry is the same in both cases.
-      element_offsets[forest->mpisize] = forest->global_num_leaf_elements;
-    }
-    else {
-      t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
-      for (i = 0; i <= mpisize; i++) {
-        element_offsets[i] = 0;
-      }
-    }
+    t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
+    element_offsets[0] = 0;
+    element_offsets[mpisize] = global_num_leaf_elements;
   }
   t8_shmem_array_end_writing (forest->element_offsets);
-
-  // In case the partition-for-coarsening flag is set, correct the partitioning if a family of elements is
-  // split across process boundaries
-  if (forest->set_for_coarsening != 0) {
-    t8_forest_pfc_correction_offsets (forest);
-  }
 }
 
 /* Find the owner of a given element.

From 211affa021a6df2117029488566014efaeaf8e2a Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 30 Dec 2025 15:33:06 +0100
Subject: [PATCH 30/35] indent

---
 src/t8_forest/t8_forest_partition.cxx | 54 ++++++++++++++++-----------
 1 file changed, 33 insertions(+), 21 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 4602cf12d3..77641eeaba 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -32,20 +32,34 @@
 #include <t8_cmesh/t8_cmesh_internal/t8_cmesh_offset.h>
 #include <t8_schemes/t8_scheme.hxx>
 
-namespace { // to prevent this implementation detail from leaking outside of this TU
-  template< typename Callable >
-  class [[maybe_unused]] DeferToScopeExit {
-   private:
-    Callable f_;
-   public:
-    DeferToScopeExit( Callable&& f ): f_{ f } {}
-    DeferToScopeExit( DeferToScopeExit const& ) = delete;
-    DeferToScopeExit( DeferToScopeExit&& ) noexcept = delete;
-    DeferToScopeExit& operator=( DeferToScopeExit const& ) = delete;
-    DeferToScopeExit& operator=( DeferToScopeExit&& ) noexcept = delete;
-    ~DeferToScopeExit(){ f_(); }
-  };
-}
+namespace
+{  // to prevent this implementation detail from leaking outside of this TU
+template <typename Callable>
+class [[maybe_unused]] DeferToScopeExit {
+ private:
+  Callable f_;
+
+ public:
+  DeferToScopeExit (Callable &&f): f_ { f }
+  {
+  }
+
+  DeferToScopeExit (DeferToScopeExit const &) = delete;
+  DeferToScopeExit (DeferToScopeExit &&) noexcept = delete;
+
+  DeferToScopeExit &
+  operator= (DeferToScopeExit const &)
+    = delete;
+  DeferToScopeExit &
+  operator= (DeferToScopeExit &&) noexcept
+    = delete;
+
+  ~DeferToScopeExit ()
+  {
+    f_ ();
+  }
+};
+}  // namespace
 
 /* We want to export the whole implementation to be callable from "C" */
 T8_EXTERN_C_BEGIN ();
@@ -519,17 +533,15 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
 
   // If the partition-for-coarsening flag is set, register the partitioning correction to be exectuted on
   // before returning. This correction prevents families of elements being split across process boundaries.
-  auto pfc = DeferToScopeExit{
-    [&](){
-      if (forest->set_for_coarsening) {
-        t8_forest_pfc_correction_offsets (forest);
-      }
+  auto pfc = DeferToScopeExit { [&] () {
+    if (forest->set_for_coarsening) {
+      t8_forest_pfc_correction_offsets (forest);
     }
-  };
+  } };
 
   // If a custom partitioning is set, all-gather the manually set element offsets.
   if (forest->set_partition_offset != 0) {
-    std::vector<t8_gloidx_t> custom_element_offsets(forest->mpisize, 0);
+    std::vector<t8_gloidx_t> custom_element_offsets (forest->mpisize, 0);
     const int retval = sc_MPI_Allgather (&forest->set_first_global_element, 1, T8_MPI_GLOIDX,
                                          custom_element_offsets.data (), 1, T8_MPI_GLOIDX, comm);
     SC_CHECK_MPI (retval);

From dfe022c7dfee30a054038b2d2a35e900920dd421 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 30 Dec 2025 15:36:05 +0100
Subject: [PATCH 31/35] fix typo

---
 src/t8_forest/t8_forest_partition.cxx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 77641eeaba..616f1d146c 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -531,7 +531,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
     return;
   }
 
-  // If the partition-for-coarsening flag is set, register the partitioning correction to be exectuted on
+  // If the partition-for-coarsening flag is set, register the partitioning correction to be executed on
   // before returning. This correction prevents families of elements being split across process boundaries.
   auto pfc = DeferToScopeExit { [&] () {
     if (forest->set_for_coarsening) {

From 25837e66bc77b4f5665b9adad4542630ef8d1725 Mon Sep 17 00:00:00 2001
From: Alex Dutka <97711898+dutkalex@users.noreply.github.com>
Date: Tue, 30 Dec 2025 15:49:26 +0100
Subject: [PATCH 32/35] revert forest commit test

---
 test/t8_forest/t8_gtest_forest_commit.cxx | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/test/t8_forest/t8_gtest_forest_commit.cxx b/test/t8_forest/t8_gtest_forest_commit.cxx
index da5b1e0857..b944fb68f1 100644
--- a/test/t8_forest/t8_gtest_forest_commit.cxx
+++ b/test/t8_forest/t8_gtest_forest_commit.cxx
@@ -129,16 +129,14 @@ t8_test_forest_commit_abp_3step (t8_forest_t forest, int maxlevel)
   t8_forest_commit (forest_balance);
 #endif
 
-  /* partition the forest with a unit weight function */
+  /* partition the forest */
   t8_forest_t forest_partition;
   t8_forest_init (&forest_partition);
-  auto weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1; };
 #if T8_TEST_LEVEL_INT < 2
   t8_forest_set_partition (forest_partition, forest_balance, 0);
 #else
   t8_forest_set_partition (forest_partition, forest_adapt, 0);
 #endif
-  t8_forest_set_partition_weights (forest_partition, weight_fcn);
   t8_forest_commit (forest_partition);
 
   return forest_partition;

From 43dbc336f867f0107f2f3cef00e5fec0ec82deb4 Mon Sep 17 00:00:00 2001
From: Thomas Spenke <thomas.spenke@dlr.de>
Date: Mon, 5 Jan 2026 12:25:56 +0100
Subject: [PATCH 33/35] Combined latest changes of dutkalex and spenke91

---
 src/t8_forest/t8_forest.cxx           |  10 +-
 src/t8_forest/t8_forest_general.h     |   4 +-
 src/t8_forest/t8_forest_partition.cxx | 206 ++++++++++++--------------
 src/t8_forest/t8_forest_types.h       |   1 +
 4 files changed, 102 insertions(+), 119 deletions(-)

diff --git a/src/t8_forest/t8_forest.cxx b/src/t8_forest/t8_forest.cxx
index 39c52f58ce..23e98768df 100644
--- a/src/t8_forest/t8_forest.cxx
+++ b/src/t8_forest/t8_forest.cxx
@@ -2704,6 +2704,8 @@ t8_forest_init (t8_forest_t *pforest)
   forest->incomplete_trees = -1;
   forest->set_partition_offset = 0;
   forest->set_first_global_element = -1;
+  forest->set_weighted_partitioning = 0;
+  forest->weight_function = nullptr;
 }
 
 int
@@ -2818,7 +2820,6 @@ t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set
   T8_ASSERT (forest->scheme == NULL);
 
   forest->set_for_coarsening = set_for_coarsening;
-  forest->weight_function = nullptr;
 
   if (set_from != NULL) {
     /* If set_from = NULL, we assume a previous forest_from was set */
@@ -2834,12 +2835,6 @@ t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set
   }
 }
 
-void
-t8_forest_set_partition_weights (t8_forest_t forest, t8_weight_fcn_t *weight_callback)
-{
-  forest->weight_function = weight_callback;
-}
-
 void
 t8_forest_set_balance (t8_forest_t forest, const t8_forest_t set_from, int no_repartition)
 {
@@ -3166,7 +3161,6 @@ t8_forest_commit (t8_forest_t forest)
           t8_forest_ref (forest->set_from);
         }
         t8_forest_set_partition (forest_partition, forest->set_from, forest->set_for_coarsening);
-        t8_forest_set_partition_weights (forest_partition, forest->weight_function);
         /* activate profiling, if this forest has profiling */
         t8_forest_set_profiling (forest_partition, forest->profile != NULL);
         /* Commit the partitioned forest */
diff --git a/src/t8_forest/t8_forest_general.h b/src/t8_forest/t8_forest_general.h
index 93611d3b10..4ac34fde53 100644
--- a/src/t8_forest/t8_forest_general.h
+++ b/src/t8_forest/t8_forest_general.h
@@ -327,8 +327,6 @@ t8_forest_get_user_function (const t8_forest_t forest);
  * false, it is not necessary to call \ref t8_forest_set_partition additionally.
  * \note This setting may not be combined with \ref t8_forest_set_copy and overwrites
  * this setting.
- * \note Calling this function sets the forest for unweighted partitioning. If weighted
- * partitioning is desired, \ref t8_forest_set_partition_weights must be called after.
  */
 void
 t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set_for_coarsening);
@@ -340,7 +338,7 @@ t8_forest_set_partition (t8_forest_t forest, const t8_forest_t set_from, int set
  * \note If \a weight_callback is null, then all the elements are assumed to have the same weight
  */
 void
-t8_forest_set_partition_weights (t8_forest_t forest, t8_weight_fcn_t *weight_callback);
+t8_forest_set_partition_weight_function (t8_forest_t forest, t8_weight_fcn_t *weight_callback);
 
 /** Set a source forest to be balanced during commit.
  * A forest is said to be balanced if each element has face neighbors of level
diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 616f1d146c..8cac381c27 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -32,35 +32,6 @@
 #include <t8_cmesh/t8_cmesh_internal/t8_cmesh_offset.h>
 #include <t8_schemes/t8_scheme.hxx>
 
-namespace
-{  // to prevent this implementation detail from leaking outside of this TU
-template <typename Callable>
-class [[maybe_unused]] DeferToScopeExit {
- private:
-  Callable f_;
-
- public:
-  DeferToScopeExit (Callable &&f): f_ { f }
-  {
-  }
-
-  DeferToScopeExit (DeferToScopeExit const &) = delete;
-  DeferToScopeExit (DeferToScopeExit &&) noexcept = delete;
-
-  DeferToScopeExit &
-  operator= (DeferToScopeExit const &)
-    = delete;
-  DeferToScopeExit &
-  operator= (DeferToScopeExit &&) noexcept
-    = delete;
-
-  ~DeferToScopeExit ()
-  {
-    f_ ();
-  }
-};
-}  // namespace
-
 /* We want to export the whole implementation to be callable from "C" */
 T8_EXTERN_C_BEGIN ();
 
@@ -501,7 +472,14 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
   }
 }
 
+void t8_forest_set_partition_weight_function(t8_forest_t forest, t8_weight_fcn_t* weight_fcn) {
+  forest->set_weighted_partitioning = 1;
+  forest->weight_function = weight_fcn;
+}
+
 // Compute forest->element_offsets according to the weight function, if provided
+/* Calculate the new element_offset for forest from
+ * the element in forest->set_from assuming a partition without element weights */
 static void
 t8_forest_partition_compute_new_offset (t8_forest_t forest)
 {
@@ -509,19 +487,20 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   T8_ASSERT (forest->set_from != NULL);
   T8_ASSERT (forest->element_offsets == NULL);
 
+  // Preparation: Define frequently used forest properties as own local variables.
   t8_forest_t forest_from = forest->set_from;
   sc_MPI_Comm comm = forest_from->mpicomm;
-  int const mpirank = forest_from->mpirank;
-  int const mpisize = forest_from->mpisize;
-  t8_gloidx_t const global_num_leaf_elements = forest_from->global_num_leaf_elements;
-  auto const weight_fcn = forest->weight_function;
+  const int mpirank = forest_from->mpirank;
+  const int mpisize = forest_from->mpisize;
+  const t8_gloidx_t global_num_leaf_elements = forest_from->global_num_leaf_elements;
+  const t8_weight_fcn_t *weight_fcn = forest->weight_function;
 
-  /* Initialize the shmem array */
+  // Initialize the shmem array.
   t8_shmem_init (comm);
   t8_shmem_set_type (comm, T8_SHMEM_BEST_TYPE);
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), mpisize + 1, comm);
 
-  // Get rid upfront of the empty forest corner case
+  // If the global numer of elements is zero, all element offsets are too and we can leave this function.
   if (global_num_leaf_elements == 0) {
     if (t8_shmem_array_start_writing (forest->element_offsets)) {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
@@ -531,103 +510,114 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
     return;
   }
 
-  // If the partition-for-coarsening flag is set, register the partitioning correction to be executed on
-  // before returning. This correction prevents families of elements being split across process boundaries.
-  auto pfc = DeferToScopeExit { [&] () {
-    if (forest->set_for_coarsening) {
-      t8_forest_pfc_correction_offsets (forest);
-    }
-  } };
+  // Define vector of manually set element offsets.
+  std::vector<t8_gloidx_t> custom_element_offsets;
 
   // If a custom partitioning is set, all-gather the manually set element offsets.
   if (forest->set_partition_offset != 0) {
-    std::vector<t8_gloidx_t> custom_element_offsets (forest->mpisize, 0);
+    custom_element_offsets.resize (forest->mpisize);
     const int retval = sc_MPI_Allgather (&forest->set_first_global_element, 1, T8_MPI_GLOIDX,
                                          custom_element_offsets.data (), 1, T8_MPI_GLOIDX, comm);
     SC_CHECK_MPI (retval);
-
-    if (t8_shmem_array_start_writing (forest->element_offsets)) {
-      t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
-      // Apply manually set element offsets.
-      std::copy_n (custom_element_offsets.data (), static_cast<size_t> (mpisize), element_offsets);
-      // The last entry is the same in both cases.
-      element_offsets[forest->mpisize] = forest->global_num_leaf_elements;
-    }
-    t8_shmem_array_end_writing (forest->element_offsets);
-    return;
   }
 
   // If no weight function is provided, compute the offsets solely based on the number of elements.
-  if (not weight_fcn) {
+  if (forest->set_weighted_partitioning == 0) {
     if (t8_shmem_array_start_writing (forest->element_offsets)) {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
-      for (int i = 0; i < mpisize; ++i) {
-        element_offsets[i] = std::floor (static_cast<double> (global_num_leaf_elements) * i / mpisize);
+
+      // Distinguish whether custom element offsets were set:
+      if (forest->set_partition_offset == 0) {
+        for (int i = 0; i < mpisize; ++i) {
+          element_offsets[i] = std::floor (static_cast<double> (global_num_leaf_elements) * i / mpisize);
+        }
       }
+      else {
+
+        // Apply manually set element offsets.
+        std::copy_n (custom_element_offsets.data (), static_cast<size_t> (mpisize), element_offsets);
+      }
+      // The last entry is the same in both cases.
       element_offsets[mpisize] = global_num_leaf_elements;
     }
     t8_shmem_array_end_writing (forest->element_offsets);
-    return;
   }
-  // Weighted load balancing:
-  // ------------------------
-  double const partition_weight = [&] () {  // sum of the weights on the local partition
-    double local_sum = 0.;
+  else {
+    // Else: Weighted load balancing:
+    // ------------------------------
+
+    // Function definition: Sum of the weights on the local partition.
+    double const partition_weight = [&] () {
+      double local_sum = 0.;
+      for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
+        for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest_from, ltreeid); ++ielm) {
+          local_sum += weight_fcn (forest_from, ltreeid, ielm);
+        }
+      }
+      return local_sum;
+    }();
+
+    // Function definition: Partial sum of the partition weights of all lower-rank processes (excluding the local rank).
+    double const partition_weight_offset = [&] () {
+      double local_offset = 0.;
+      double local_partition_weight = partition_weight;  // because MPI does not like const variables
+      sc_MPI_Exscan (&local_partition_weight, &local_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, comm);
+      return mpirank > 0 ? local_offset : 0;  // because the result of MPI_Exscan is undefined on rank 0
+    }();
+
+    // Function definition: Complete sum of the partition weights.
+    double const forest_weight = [&] () {
+      double total_weight = partition_weight_offset + partition_weight;
+      sc_MPI_Bcast (&total_weight, 1, sc_MPI_DOUBLE, mpisize - 1, comm);
+      return total_weight;
+    }();
+
+    // The [rank_begin, rank_end) slice of the new offsets will land in the local partition.
+    int const rank_begin = std::ceil (mpisize * partition_weight_offset / forest_weight);
+    int const rank_end = std::ceil (mpisize * (partition_weight_offset + partition_weight) / forest_weight);
+    std::vector<t8_gloidx_t> local_offsets (rank_end - rank_begin, 0);
+
+    // Prepare compuation of local element offsets.
+    double accumulated_weight = partition_weight_offset;
+    t8_gloidx_t global_elm_idx = t8_forest_get_first_local_leaf_element_id (forest_from);
+    int iproc = rank_begin;
+
+    // Loop over local trees and their elements:
+    // - Add element weight to accumulated sum.
+    // - If accumulated weight is above ideal bound for proc iproc, set local offset and go to next process.
     for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
       for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest_from, ltreeid); ++ielm) {
-        local_sum += weight_fcn (forest_from, ltreeid, ielm);
+        T8_ASSERT (0 <= global_elm_idx && global_elm_idx < global_num_leaf_elements);
+        accumulated_weight += weight_fcn (forest_from, ltreeid, ielm);
+        while (accumulated_weight
+              > forest_weight * iproc / mpisize) {  // there may be empty partitions, hence while and not if
+          T8_ASSERT (rank_begin <= iproc && iproc < rank_end);
+          local_offsets[iproc - rank_begin] = global_elm_idx;
+          ++iproc;
+        }
+        ++global_elm_idx;
       }
     }
-    return local_sum;
-  }();
-
-  double const partition_weight_offset
-    = [&] () {  // partial sum of the partition weights of all lower-rank processes (excluding the local rank)
-        double local_offset = 0.;
-        double local_partition_weight = partition_weight;  // because MPI does not like const variables
-        sc_MPI_Exscan (&local_partition_weight, &local_offset, 1, sc_MPI_DOUBLE, sc_MPI_SUM, comm);
-        return mpirank > 0 ? local_offset : 0;  // because the result of MPI_Exscan is undefined on rank 0
-      }();
-
-  double const forest_weight = [&] () {  // complete sum of the partition weights
-    double total_weight = partition_weight_offset + partition_weight;
-    sc_MPI_Bcast (&total_weight, 1, sc_MPI_DOUBLE, mpisize - 1, comm);
-    return total_weight;
-  }();
-
-  // The [rank_begin, rank_end) slice of the new offsets land in the local partition
-  int const rank_begin = std::ceil (mpisize * partition_weight_offset / forest_weight);
-  int const rank_end = std::ceil (mpisize * (partition_weight_offset + partition_weight) / forest_weight);
-  std::vector<t8_gloidx_t> local_offsets (rank_end - rank_begin, 0);
-
-  double accumulated_weight = partition_weight_offset;
-  t8_gloidx_t global_elm_idx = t8_forest_get_first_local_leaf_element_id (forest_from);
-  int i = rank_begin;
-
-  for (t8_locidx_t ltreeid = 0; ltreeid < t8_forest_get_num_local_trees (forest_from); ++ltreeid) {
-    for (t8_locidx_t ielm = 0; ielm < t8_forest_get_tree_num_leaf_elements (forest_from, ltreeid); ++ielm) {
-      T8_ASSERT (0 <= global_elm_idx && global_elm_idx < global_num_leaf_elements);
-      accumulated_weight += weight_fcn (forest_from, ltreeid, ielm);
-      while (accumulated_weight
-             > forest_weight * i / mpisize) {  // there may be empty partitions, hence while and not if
-        T8_ASSERT (rank_begin <= i && i < rank_end);
-        local_offsets[i - rank_begin] = global_elm_idx;
-        ++i;
-      }
-      ++global_elm_idx;
+    T8_ASSERT (iproc == rank_end);  // i.e. local_offsets has been filled properly
+
+    // Allgather local offsets into the global element offset array.
+    t8_shmem_array_allgatherv (local_offsets.data (), local_offsets.size (), T8_MPI_GLOIDX, forest->element_offsets,
+                              T8_MPI_GLOIDX, comm);
+
+    // Add first and last entry to element offsets.
+    if (t8_shmem_array_start_writing (forest->element_offsets)) {
+      t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
+      element_offsets[0] = 0;
+      element_offsets[mpisize] = global_num_leaf_elements;
     }
+    t8_shmem_array_end_writing (forest->element_offsets);
   }
-  T8_ASSERT (i == rank_end);  // i.e. local_offsets has been filled properly
-
-  t8_shmem_array_allgatherv (local_offsets.data (), local_offsets.size (), T8_MPI_GLOIDX, forest->element_offsets,
-                             T8_MPI_GLOIDX, comm);
 
-  if (t8_shmem_array_start_writing (forest->element_offsets)) {
-    t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
-    element_offsets[0] = 0;
-    element_offsets[mpisize] = global_num_leaf_elements;
+  // In case the partition-for-coarsening flag is set, correct the partitioning if a family of elements is
+  // split across process boundaries.
+  if (forest->set_for_coarsening != 0) {
+    t8_forest_pfc_correction_offsets (forest);
   }
-  t8_shmem_array_end_writing (forest->element_offsets);
 }
 
 /* Find the owner of a given element.
diff --git a/src/t8_forest/t8_forest_types.h b/src/t8_forest/t8_forest_types.h
index a855e4adcd..02aa51277b 100644
--- a/src/t8_forest/t8_forest_types.h
+++ b/src/t8_forest/t8_forest_types.h
@@ -76,6 +76,7 @@ typedef struct t8_forest
   int set_level;                    /**< Level to use in new construction. */
   int set_for_coarsening;           /**< Change partition to allow
                                                      for one round of coarsening */
+  int set_weighted_partitioning;    /**< Flag indicating whether a weighting function is used for partitioning.*/
   t8_weight_fcn_t *weight_function; /**< Pointer to user defined element weight function. Can be null. */
 
   sc_MPI_Comm mpicomm; /**< MPI communicator to use. */

From cb568a26f1bbdf001506ac4bb996ba763608187e Mon Sep 17 00:00:00 2001
From: Thomas Spenke <thomas.spenke@dlr.de>
Date: Mon, 5 Jan 2026 12:29:05 +0100
Subject: [PATCH 34/35] Added test case for weighted partitioning

---
 .../t8_gtest_weighted_partitioning.cxx        | 231 ++++++++++++++++++
 1 file changed, 231 insertions(+)
 create mode 100644 test/t8_forest/t8_gtest_weighted_partitioning.cxx

diff --git a/test/t8_forest/t8_gtest_weighted_partitioning.cxx b/test/t8_forest/t8_gtest_weighted_partitioning.cxx
new file mode 100644
index 0000000000..c1d25fa388
--- /dev/null
+++ b/test/t8_forest/t8_gtest_weighted_partitioning.cxx
@@ -0,0 +1,231 @@
+/*
+  This file is part of t8code.
+  t8code is a C library to manage a collection (a forest) of multiple
+  connected adaptive space-trees of general element classes in parallel.
+
+  Copyright (C) 2025 the developers
+
+  t8code is free software; you can redistribute it and/or modify
+  it under the terms of the GNU General Public License as published by
+  the Free Software Foundation; either version 2 of the License, or
+  (at your option) any later version.
+
+  t8code is distributed in the hope that it will be useful,
+  but WITHOUT ANY WARRANTY; without even the implied warranty of
+  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+  GNU General Public License for more details.
+
+  You should have received a copy of the GNU General Public License
+  along with t8code; if not, write to the Free Software Foundation, Inc.,
+  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+*/
+
+/**
+ * \file In this file, we test the weighted-partitioning feature of t8code.
+ *
+ * It allows defining a weighting function that assigns a weight to each element
+ * in the forest. The partitioning will then try to establish a (as far as possible)
+ * ideal balancing of the weights, rather than the element numbers as "usual".
+ *
+ * The following steps are performed per example cmesh and scheme, with sanity
+ * checks validating each of them:
+ *
+ * (1.) A uniform, partitioned base forest is created.
+ *
+ * (2.) In a first test, the weighting function is equal to one for all elements.
+ *      This allows a simple sanity check: The resulting partition has to be the
+ *      same as without assigning weights.
+ *
+ * (3.) Next, we test a non-uniform weight function. As an example, we defined the
+ *      weight function such that an element's weight is simply its global element id.
+ *      This allows to "manually" compute the expected element offsets to then compare 
+ *      them.
+ */
+
+// testing
+#include <gtest/gtest.h>
+#include <test/t8_gtest_schemes.hxx>
+#include "test/t8_cmesh_generator/t8_cmesh_example_sets.hxx"
+#include <test/t8_gtest_macros.hxx>
+#include <test/t8_gtest_custom_assertion.hxx>
+
+// t8code
+#include <t8.h>
+#include <t8_cmesh/t8_cmesh.h>
+#include <t8_forest/t8_forest_general.h>
+#include <t8_forest/t8_forest_partition.h>
+#include <t8_forest/t8_forest_types.h>
+
+/**
+ * Test class for validating t8code's weighted partitioning. See file description for details on the testing strategy.
+*/
+class t8_test_weighted_partitioning_test: public testing::TestWithParam<std::tuple<int, cmesh_example_base *>> {
+
+ protected:
+  /**
+   * Setup routine of the test suite.
+  */
+  void
+  SetUp () override
+  {
+    // Get scheme.
+    const int scheme_id = std::get<0> (GetParam ());
+    scheme = create_from_scheme_id (scheme_id);
+
+    // Construct cmesh.
+    cmesh = std::get<1> (GetParam ())->cmesh_create ();
+
+    // Skip empty meshes.
+    if (t8_cmesh_is_empty (cmesh)) {
+      GTEST_SKIP ();
+    }
+  }
+
+  /**
+   * Tear down routine of the test suite.
+   *
+   * Makes sure to unref the cmesh and the scheme.
+  */
+  void
+  TearDown () override
+  {
+    t8_cmesh_destroy (&cmesh);
+    scheme->unref ();
+  }
+
+  // Member variables:
+  const t8_scheme *scheme;
+  t8_cmesh_t cmesh;
+};
+
+/**
+ * Main function of testing suite. See doxygen file description for details on the testing strategy.
+*/
+TEST_P (t8_test_weighted_partitioning_test, test_weighted_partitioning)
+{
+
+  // -------------------------------------------
+  // ----- (1.) Create uniform base forest -----
+  // -------------------------------------------
+
+  t8_debugf ("Create uniform base forest.\n");
+
+  // Initial uniform refinement level
+  const int level = 2;
+
+  // Increase reference counters of cmesh and scheme to avoid reaching zero.
+  t8_cmesh_ref (cmesh);
+  scheme->ref ();
+
+  // Create initial, uniform base forest.
+  t8_forest_t base_forest = t8_forest_new_uniform (cmesh, scheme, level, 0, sc_MPI_COMM_WORLD);
+
+  // ---------------------------------------------
+  // ----- (2.) Test 1: All weights the same -----
+  // ---------------------------------------------
+  t8_debugf ("Create forest with uniform partition weights.\n");
+
+  // Initialize forest
+  t8_forest_t forest_uniform_weights;
+  t8_forest_init (&forest_uniform_weights);
+
+  // Define uniform weight fcn.
+  t8_weight_fcn_t* uniform_weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1; };
+
+  // Prepare partitioning.
+  t8_forest_set_partition(forest_uniform_weights, base_forest, 0);
+  t8_forest_set_partition_weight_function(forest_uniform_weights, uniform_weight_fcn);
+
+  // Commit forest (without destroying base_forest).
+  t8_forest_ref (base_forest);
+  t8_forest_commit (forest_uniform_weights);
+
+  // Create partitioned forest without weights for comparison.
+  t8_forest_t forest_no_weights;
+  t8_forest_init (&forest_no_weights);
+  t8_forest_set_partition(forest_no_weights, base_forest, 0);
+  t8_forest_ref (base_forest);
+  t8_forest_commit (forest_no_weights);
+
+  // Make sure they are the same.
+  EXPECT_FOREST_EQ (forest_uniform_weights, forest_no_weights);
+
+  // ---------------------------------------------
+  // ----- (3.) Test 2: Non-uniform weights ------
+  // ---------------------------------------------
+  t8_debugf ("Create forest with non-uniform partition weights.\n");
+
+  // (3a.) Apply weighted partitioning:
+  // ----------------------------------
+
+  // Define non-uniform weight fcn: For this test, we pick the element number as element weight.
+  t8_weight_fcn_t* weight_fcn = [] (t8_forest_t forest, t8_locidx_t ltree_id, t8_locidx_t ele_in_tree) -> double {
+    const t8_gloidx_t gelem_id = t8_forest_get_first_local_leaf_element_id(forest) + t8_forest_get_tree_element_offset(forest, ltree_id) + ele_in_tree;
+    return gelem_id; };
+
+  // Prepare forest.
+  t8_forest_t forest_weighted;
+  t8_forest_init (&forest_weighted);
+  t8_forest_set_partition(forest_weighted, base_forest, 0);
+  t8_forest_set_partition_weight_function(forest_weighted, weight_fcn);
+
+  // Commit forest (without destroying base_forest).
+  t8_forest_ref (base_forest);
+  t8_forest_commit (forest_weighted);
+
+  // (3b.) Determine expected partitioning:
+  // --------------------------------------
+
+  // Preparation: MPI variablies and global number of leaf elements.
+  const int mpirank = base_forest->mpirank;
+  const int mpisize = base_forest->mpisize;
+  const t8_gloidx_t num_global_leaves = t8_forest_get_global_num_leaf_elements(base_forest);
+
+  // Total sum of the weights for this case (Gauss sum formula for num_global_leaves - 1)
+  const double total_weight_sum = num_global_leaves * (num_global_leaves - 1) / 2;
+
+  // "Manually" compute the element offsets expected for this distribution.
+  // Note: For simplicty, every process computes all offsets here, so no MPI communication is required.
+  std::vector<t8_gloidx_t> element_offsets(mpisize+1);
+  double sum = 0.0;
+  int on_rank = 0;
+  element_offsets[0] = 0;
+  for(t8_gloidx_t ielem = 0; ielem < num_global_leaves; ielem++) {
+
+    // Add weight (which matches element number here)
+    sum += ielem;
+
+    // If we are above ideal bound for current process, add offset to vector and go to next rank.
+    if(sum > total_weight_sum * (on_rank+1) / mpisize) {
+      on_rank ++;
+      std::fill(element_offsets.begin()+on_rank, element_offsets.end(), ielem);
+      t8_global_productionf("Computed offset for rank %i: %li \n", on_rank, element_offsets[on_rank]);
+    }
+  }
+  // The last entry is always the global number of elements.
+  element_offsets[mpisize] = num_global_leaves;
+
+  // Determine number of local elements.
+  t8_locidx_t num_local_leaves = element_offsets[mpirank+1] - element_offsets[mpirank];
+
+  // (3c.) Comparison:
+  // -----------------
+
+  // Make sure the number of local leaves match.
+  EXPECT_EQ( t8_forest_get_local_num_leaf_elements(forest_weighted), num_local_leaves);
+
+  // ---------------------------
+  // ----- Memory clean-up -----
+  // ---------------------------
+
+  // Destroy the forests.
+  t8_forest_unref (&base_forest);
+  t8_forest_unref (&forest_uniform_weights);
+  t8_forest_unref (&forest_no_weights);
+  t8_forest_unref (&forest_weighted);
+
+}
+
+// Instantiate parameterized test to be run for all schemes and example cmeshes.
+INSTANTIATE_TEST_SUITE_P (t8_gtest_weighted_partitioning, t8_test_weighted_partitioning_test,
+                          testing::Combine (AllSchemeCollections, AllCmeshsParam), pretty_print_base_example_scheme);

From 7446ce771c00007e794c97d920f3fbb313d010e1 Mon Sep 17 00:00:00 2001
From: Thomas Spenke <thomas.spenke@dlr.de>
Date: Mon, 5 Jan 2026 12:58:20 +0100
Subject: [PATCH 35/35] Fix typos and indentation

---
 src/t8_forest/t8_forest_partition.cxx         | 12 +++---
 src/t8_netcdf.c                               |  2 +-
 .../t8_gtest_weighted_partitioning.cxx        | 43 ++++++++++---------
 3 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/src/t8_forest/t8_forest_partition.cxx b/src/t8_forest/t8_forest_partition.cxx
index 8cac381c27..04341a3f70 100644
--- a/src/t8_forest/t8_forest_partition.cxx
+++ b/src/t8_forest/t8_forest_partition.cxx
@@ -472,7 +472,9 @@ t8_forest_partition_create_tree_offsets (t8_forest_t forest)
   }
 }
 
-void t8_forest_set_partition_weight_function(t8_forest_t forest, t8_weight_fcn_t* weight_fcn) {
+void
+t8_forest_set_partition_weight_function (t8_forest_t forest, t8_weight_fcn_t *weight_fcn)
+{
   forest->set_weighted_partitioning = 1;
   forest->weight_function = weight_fcn;
 }
@@ -500,7 +502,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
   t8_shmem_set_type (comm, T8_SHMEM_BEST_TYPE);
   t8_shmem_array_init (&forest->element_offsets, sizeof (t8_gloidx_t), mpisize + 1, comm);
 
-  // If the global numer of elements is zero, all element offsets are too and we can leave this function.
+  // If the global number of elements is zero, all element offsets are too and we can leave this function.
   if (global_num_leaf_elements == 0) {
     if (t8_shmem_array_start_writing (forest->element_offsets)) {
       t8_gloidx_t *element_offsets = t8_shmem_array_get_gloidx_array_for_writing (forest->element_offsets);
@@ -577,7 +579,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
     int const rank_end = std::ceil (mpisize * (partition_weight_offset + partition_weight) / forest_weight);
     std::vector<t8_gloidx_t> local_offsets (rank_end - rank_begin, 0);
 
-    // Prepare compuation of local element offsets.
+    // Prepare computation of local element offsets.
     double accumulated_weight = partition_weight_offset;
     t8_gloidx_t global_elm_idx = t8_forest_get_first_local_leaf_element_id (forest_from);
     int iproc = rank_begin;
@@ -590,7 +592,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
         T8_ASSERT (0 <= global_elm_idx && global_elm_idx < global_num_leaf_elements);
         accumulated_weight += weight_fcn (forest_from, ltreeid, ielm);
         while (accumulated_weight
-              > forest_weight * iproc / mpisize) {  // there may be empty partitions, hence while and not if
+               > forest_weight * iproc / mpisize) {  // there may be empty partitions, hence while and not if
           T8_ASSERT (rank_begin <= iproc && iproc < rank_end);
           local_offsets[iproc - rank_begin] = global_elm_idx;
           ++iproc;
@@ -602,7 +604,7 @@ t8_forest_partition_compute_new_offset (t8_forest_t forest)
 
     // Allgather local offsets into the global element offset array.
     t8_shmem_array_allgatherv (local_offsets.data (), local_offsets.size (), T8_MPI_GLOIDX, forest->element_offsets,
-                              T8_MPI_GLOIDX, comm);
+                               T8_MPI_GLOIDX, comm);
 
     // Add first and last entry to element offsets.
     if (t8_shmem_array_start_writing (forest->element_offsets)) {
diff --git a/src/t8_netcdf.c b/src/t8_netcdf.c
index faf8a34423..2b5c74fa23 100644
--- a/src/t8_netcdf.c
+++ b/src/t8_netcdf.c
@@ -46,7 +46,7 @@ t8_netcdf_create_integer_var (const char *var_name, const char *var_long_name, c
                               sc_array_t *var_data)
 {
   t8_netcdf_variable_type_t var_type;
-  /* Check whether 32-bit (4-byte) integer or 64-bit integer data sholud be written */
+  /* Check whether 32-bit (4-byte) integer or 64-bit integer data should be written */
   var_type = (var_data->elem_size > 4) ? T8_NETCDF_INT64 : T8_NETCDF_INT;
   return t8_netcdf_create_var (var_type, var_name, var_long_name, var_unit, var_data);
 }
diff --git a/test/t8_forest/t8_gtest_weighted_partitioning.cxx b/test/t8_forest/t8_gtest_weighted_partitioning.cxx
index c1d25fa388..5e178d2c6d 100644
--- a/test/t8_forest/t8_gtest_weighted_partitioning.cxx
+++ b/test/t8_forest/t8_gtest_weighted_partitioning.cxx
@@ -130,11 +130,11 @@ TEST_P (t8_test_weighted_partitioning_test, test_weighted_partitioning)
   t8_forest_init (&forest_uniform_weights);
 
   // Define uniform weight fcn.
-  t8_weight_fcn_t* uniform_weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1; };
+  t8_weight_fcn_t *uniform_weight_fcn = [] (t8_forest_t, t8_locidx_t, t8_locidx_t) -> double { return 1; };
 
   // Prepare partitioning.
-  t8_forest_set_partition(forest_uniform_weights, base_forest, 0);
-  t8_forest_set_partition_weight_function(forest_uniform_weights, uniform_weight_fcn);
+  t8_forest_set_partition (forest_uniform_weights, base_forest, 0);
+  t8_forest_set_partition_weight_function (forest_uniform_weights, uniform_weight_fcn);
 
   // Commit forest (without destroying base_forest).
   t8_forest_ref (base_forest);
@@ -143,7 +143,7 @@ TEST_P (t8_test_weighted_partitioning_test, test_weighted_partitioning)
   // Create partitioned forest without weights for comparison.
   t8_forest_t forest_no_weights;
   t8_forest_init (&forest_no_weights);
-  t8_forest_set_partition(forest_no_weights, base_forest, 0);
+  t8_forest_set_partition (forest_no_weights, base_forest, 0);
   t8_forest_ref (base_forest);
   t8_forest_commit (forest_no_weights);
 
@@ -159,15 +159,17 @@ TEST_P (t8_test_weighted_partitioning_test, test_weighted_partitioning)
   // ----------------------------------
 
   // Define non-uniform weight fcn: For this test, we pick the element number as element weight.
-  t8_weight_fcn_t* weight_fcn = [] (t8_forest_t forest, t8_locidx_t ltree_id, t8_locidx_t ele_in_tree) -> double {
-    const t8_gloidx_t gelem_id = t8_forest_get_first_local_leaf_element_id(forest) + t8_forest_get_tree_element_offset(forest, ltree_id) + ele_in_tree;
-    return gelem_id; };
+  t8_weight_fcn_t *weight_fcn = [] (t8_forest_t forest, t8_locidx_t ltree_id, t8_locidx_t ele_in_tree) -> double {
+    const t8_gloidx_t gelem_id = t8_forest_get_first_local_leaf_element_id (forest)
+                                 + t8_forest_get_tree_element_offset (forest, ltree_id) + ele_in_tree;
+    return gelem_id;
+  };
 
   // Prepare forest.
   t8_forest_t forest_weighted;
   t8_forest_init (&forest_weighted);
-  t8_forest_set_partition(forest_weighted, base_forest, 0);
-  t8_forest_set_partition_weight_function(forest_weighted, weight_fcn);
+  t8_forest_set_partition (forest_weighted, base_forest, 0);
+  t8_forest_set_partition_weight_function (forest_weighted, weight_fcn);
 
   // Commit forest (without destroying base_forest).
   t8_forest_ref (base_forest);
@@ -176,43 +178,43 @@ TEST_P (t8_test_weighted_partitioning_test, test_weighted_partitioning)
   // (3b.) Determine expected partitioning:
   // --------------------------------------
 
-  // Preparation: MPI variablies and global number of leaf elements.
+  // Preparation: MPI variables and global number of leaf elements.
   const int mpirank = base_forest->mpirank;
   const int mpisize = base_forest->mpisize;
-  const t8_gloidx_t num_global_leaves = t8_forest_get_global_num_leaf_elements(base_forest);
+  const t8_gloidx_t num_global_leaves = t8_forest_get_global_num_leaf_elements (base_forest);
 
   // Total sum of the weights for this case (Gauss sum formula for num_global_leaves - 1)
   const double total_weight_sum = num_global_leaves * (num_global_leaves - 1) / 2;
 
   // "Manually" compute the element offsets expected for this distribution.
-  // Note: For simplicty, every process computes all offsets here, so no MPI communication is required.
-  std::vector<t8_gloidx_t> element_offsets(mpisize+1);
+  // Note: For simplicity, every process computes all offsets here, so no MPI communication is required.
+  std::vector<t8_gloidx_t> element_offsets (mpisize + 1);
   double sum = 0.0;
   int on_rank = 0;
   element_offsets[0] = 0;
-  for(t8_gloidx_t ielem = 0; ielem < num_global_leaves; ielem++) {
+  for (t8_gloidx_t ielem = 0; ielem < num_global_leaves; ielem++) {
 
     // Add weight (which matches element number here)
     sum += ielem;
 
     // If we are above ideal bound for current process, add offset to vector and go to next rank.
-    if(sum > total_weight_sum * (on_rank+1) / mpisize) {
-      on_rank ++;
-      std::fill(element_offsets.begin()+on_rank, element_offsets.end(), ielem);
-      t8_global_productionf("Computed offset for rank %i: %li \n", on_rank, element_offsets[on_rank]);
+    if (sum > total_weight_sum * (on_rank + 1) / mpisize) {
+      on_rank++;
+      std::fill (element_offsets.begin () + on_rank, element_offsets.end (), ielem);
+      t8_global_productionf ("Computed offset for rank %i: %li \n", on_rank, element_offsets[on_rank]);
     }
   }
   // The last entry is always the global number of elements.
   element_offsets[mpisize] = num_global_leaves;
 
   // Determine number of local elements.
-  t8_locidx_t num_local_leaves = element_offsets[mpirank+1] - element_offsets[mpirank];
+  t8_locidx_t num_local_leaves = element_offsets[mpirank + 1] - element_offsets[mpirank];
 
   // (3c.) Comparison:
   // -----------------
 
   // Make sure the number of local leaves match.
-  EXPECT_EQ( t8_forest_get_local_num_leaf_elements(forest_weighted), num_local_leaves);
+  EXPECT_EQ (t8_forest_get_local_num_leaf_elements (forest_weighted), num_local_leaves);
 
   // ---------------------------
   // ----- Memory clean-up -----
@@ -223,7 +225,6 @@ TEST_P (t8_test_weighted_partitioning_test, test_weighted_partitioning)
   t8_forest_unref (&forest_uniform_weights);
   t8_forest_unref (&forest_no_weights);
   t8_forest_unref (&forest_weighted);
-
 }
 
 // Instantiate parameterized test to be run for all schemes and example cmeshes.