Feature: extend the partition algorithm to handle weights

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);

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);
   }

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);
+  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);
 

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);
   }

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);
   }

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);
   }

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);
   }

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 */
@@ -1952,7 +1952,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);
@@ -1974,7 +1974,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);
@@ -2315,7 +2315,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)
@@ -2798,7 +2798,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_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);
@@ -2808,6 +2809,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 */
@@ -2956,9 +2958,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
  */
@@ -3148,7 +3150,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);
+        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 */
@@ -3169,7 +3172,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) {

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) {