Feature: Add Partition and Balance to mesh handle

mesh_handle/adapt.hxx

@@ -120,8 +120,8 @@ class AdaptRegistry {
  public:
   /** Static function to register \a context using \a forest as identifier. 
    * This makes the context publicly available using the Registry.
-   * \param [in] forest  The forest identifier. In our case, this is the forest to be adapted 
-   *                     and not the forest from which we adapt.
+   * \param [in] forest  The forest identifier. In our case, this is the forest from which we adapt the mesh because 
+   *                we do not change this forest during adaptation such that it is valid unique identifier.
    * \param [in] context The context to register. Use unique pointer to ensure proper memory management and ownership.
    */
   static void
@@ -135,7 +135,7 @@ class AdaptRegistry {
   }
 
   /** Static function to unregister a context using \a forest as identifier. 
-   * \param [in] forest The forest identifier. In our case, this is the forest to be adapted.
+   * \param [in] forest The forest identifier. In our case, this is the forest from which we adapt.
    */
   static void
   unregister_context (t8_forest_t forest)
@@ -146,7 +146,7 @@ class AdaptRegistry {
   }
 
   /** Getter for a context using \a forest as identifier. 
-   * \param [in] forest The forest identifier. In our case, this is the forest to be adapted.
+   * \param [in] forest The forest identifier. In our case, this is the forest from which we adapt.
    * \return Pointer to the context registered with the id \a forest if found, nullptr otherwise.
    */
   static MeshAdaptContextBase*
@@ -193,7 +193,7 @@ mesh_adapt_callback_wrapper ([[maybe_unused]] t8_forest_t forest, t8_forest_t fo
 {
   // Get static adapt context from the registry.
   // Via this, we can access the mesh handle and the user defined adapt callback that uses mesh handle functionality.
-  auto* context = AdaptRegistry::get (forest);
+  auto* context = AdaptRegistry::get (forest_from);
   if (!context) {
     t8_global_infof (
       "Something went wrong while registering the adaptation callbacks. Please check your implementation.");

mesh_handle/mesh.hxx

@@ -30,6 +30,8 @@
 #include "element.hxx"
 #include "competence_pack.hxx"
 #include "adapt.hxx"
+#include "t8_forest/t8_forest_balance.h"
+#include "t8_forest/t8_forest_types.h"
 #include <t8_forest/t8_forest_general.h>
 #include <t8_forest/t8_forest_ghost.h>
 #include <vector>
@@ -162,6 +164,17 @@ class mesh {
     return m_forest;
   }
 
+  /** Check if the local elements of the mesh are balanced. 
+  * The mesh is said to be balanced if each element has face neighbors of level
+  * at most +1 or -1 of the element's level.
+  * \return true if the local elements are balanced, false otherwise.
+  */
+  bool
+  is_balanced ()
+  {
+    return t8_forest_is_balanced (m_forest);
+  }
+
   // --- Methods to access elements. ---
   /**
    * Returns a constant iterator to the first (local) mesh element.
@@ -270,19 +283,62 @@ class mesh {
     }
     // Create and register adaptation context holding the mesh handle and the user defined callback.
     detail::AdaptRegistry::register_context (
-      m_uncommitted_forest.value (),
-      std::make_unique<detail::MeshAdaptContext<SelfType>> (*this, std::move (adapt_callback)));
+      m_forest, std::make_unique<detail::MeshAdaptContext<SelfType>> (*this, std::move (adapt_callback)));
 
     // Set up the forest for adaptation using the wrapper callback.
     t8_forest_set_adapt (m_uncommitted_forest.value (), m_forest, detail::mesh_adapt_callback_wrapper, recursive);
   }
 
+  /** If this function is called, the mesh will be partitioned on committing.
+   * The partitioning is done according to the SFC and each rank is assigned
+   * the same (maybe +1) number of elements.
+   * \note The partition is carried out only when \ref commit is called.
+   * \note This setting can be combined with \ref set_adapt and \ref set_balance. The order in which
+   * these operations are executed is always 1) Adapt 2) Partition 3) Balance.
+   * \param [in] set_for_coarsening If true, the partitions are choose such that coarsening 
+   *        an element once is a process local operation. Default is false.
+   */
+  void
+  set_partition (bool set_for_coarsening = false)
+  {
+    if (!m_uncommitted_forest.has_value ()) {
+      t8_forest_t new_forest;
+      t8_forest_init (&new_forest);
+      m_uncommitted_forest = new_forest;
+    }
+    t8_forest_set_partition (m_uncommitted_forest.value (), m_forest, set_for_coarsening);
+  }
+
+  /** If this function is called, the mesh will be balanced on committing.
+   * The mesh is said to be balanced if each element has face neighbors of level
+   * at most +1 or -1 of the element's level.
+   * \note The balance is carried out only when \ref commit is called.
+   * \param [in] no_repartition Balance constructs several intermediate steps that
+   *       are refined from each other. In order to maintain a balanced load, a repartitioning is performed in each 
+   *       round and the resulting mesh is load-balanced per default. 
+   *       Set \a no_repartition to true if this behaviour is not desired.
+   *       If \a no_repartition is false (default), an additional call of \ref set_partition is not necessary.
+   * \note This setting can be combined with \ref set_adapt and \ref set_partition. The order in which
+   * these operations are executed is always 1) Adapt 2) Partition 3) Balance.
+   */
+  void
+  set_balance (bool no_repartition = false)
+  {
+    if (!m_uncommitted_forest.has_value ()) {
+      t8_forest_t new_forest;
+      t8_forest_init (&new_forest);
+      m_uncommitted_forest = new_forest;
+    }
+    // Disable repartitioning and let the user call set_partition if desired.
+    t8_forest_set_balance (m_uncommitted_forest.value (), m_forest, no_repartition);
+  }
+
   /** Enable or disable the creation of a layer of ghost elements.
-  * \param [in]      do_ghost  If true a ghost layer will be created.
-  * \param [in]      ghost_type Controls which neighbors count as ghost elements,
-  *                             currently only T8_GHOST_FACES is supported. This value
-  *                             is ignored if \a do_ghost = false.
-  */
+   * \param [in]      do_ghost  If true a ghost layer will be created.
+   * \param [in]      ghost_type Controls which neighbors count as ghost elements,
+   *                             currently only T8_GHOST_FACES is supported. This value
+   *                             is ignored if \a do_ghost = false.
+   */
   void
   set_ghost (bool do_ghost = true, t8_ghost_type_t ghost_type = T8_GHOST_FACES)
   {
@@ -313,16 +369,18 @@ class mesh {
     if (m_uncommitted_forest.value ()->set_from == NULL) {
       t8_forest_set_copy (m_uncommitted_forest.value (), m_forest);
     }
+    t8_forest_ref (m_forest);
     t8_forest_commit (m_uncommitted_forest.value ());
     // Check if we adapted and unregister the adapt context if so.
     if (detail::AdaptRegistry::get (m_uncommitted_forest.value ()) != nullptr) {
-      detail::AdaptRegistry::unregister_context (m_uncommitted_forest.value ());
+      detail::AdaptRegistry::unregister_context (m_forest);
       if (!std::is_void<TElementDataType>::value) {
         t8_global_infof (
           "Please note that the element data is not interpolated automatically during adaptation. Use the "
           "function set_element_data() to provide new adapted element data.\n");
       }
     }
+    t8_forest_unref (&m_forest);
     // Update underlying forest of the mesh.
     m_forest = m_uncommitted_forest.value ();
     m_uncommitted_forest = std::nullopt;

test/CMakeLists.txt

@@ -209,13 +209,13 @@ endif()
 add_t8_cpp_test( NAME t8_gtest_vector_split_serial          SOURCES t8_helper_functions/t8_gtest_vector_split.cxx )
 
 if( T8CODE_BUILD_MESH_HANDLE )
-  add_t8_cpp_test( NAME t8_gtest_mesh_handle_parallel             SOURCES mesh_handle/t8_gtest_mesh_handle.cxx )
-  add_t8_cpp_test( NAME t8_gtest_compare_handle_to_forest_serial  SOURCES mesh_handle/t8_gtest_compare_handle_to_forest.cxx )
-  add_t8_cpp_test( NAME t8_gtest_handle_ghost_parallel            SOURCES mesh_handle/t8_gtest_ghost.cxx )
-  add_t8_cpp_test( NAME t8_gtest_custom_competence_serial         SOURCES mesh_handle/t8_gtest_custom_competence.cxx )
-  add_t8_cpp_test( NAME t8_gtest_cache_competence_serial          SOURCES mesh_handle/t8_gtest_cache_competence.cxx )
-  add_t8_cpp_test( NAME t8_gtest_handle_data_parallel             SOURCES mesh_handle/t8_gtest_handle_data.cxx )
-  add_t8_cpp_test( NAME t8_gtest_handle_adapt_serial             SOURCES mesh_handle/t8_gtest_adapt.cxx )
+  add_t8_cpp_test( NAME t8_gtest_mesh_handle_parallel               SOURCES mesh_handle/t8_gtest_mesh_handle.cxx )
+  add_t8_cpp_test( NAME t8_gtest_compare_handle_to_forest_serial    SOURCES mesh_handle/t8_gtest_compare_handle_to_forest.cxx )
+  add_t8_cpp_test( NAME t8_gtest_handle_ghost_parallel              SOURCES mesh_handle/t8_gtest_ghost.cxx )
+  add_t8_cpp_test( NAME t8_gtest_custom_competence_serial           SOURCES mesh_handle/t8_gtest_custom_competence.cxx )
+  add_t8_cpp_test( NAME t8_gtest_cache_competence_serial            SOURCES mesh_handle/t8_gtest_cache_competence.cxx )
+  add_t8_cpp_test( NAME t8_gtest_handle_data_parallel               SOURCES mesh_handle/t8_gtest_handle_data.cxx )
+  add_t8_cpp_test( NAME t8_gtest_adapt_partition_balance_parallel   SOURCES mesh_handle/t8_gtest_adapt_partition_balance.cxx )
 endif()
 
 copy_test_file( test_cube_unstructured_1.inp )

test/mesh_handle/t8_gtest_adapt.cxx → test/mesh_handle/t8_gtest_adapt_partition_balance.cxx

@@ -21,11 +21,13 @@ along with t8code; if not, write to the Free Software Foundation, Inc.,
 */
 
 /**
- * \file t8_gtest_adapt.cxx
- * Tests for the adapt routines of mesh handles. 
- * This tests uses the callback and user data of tutorial step 3 as example.
+ * \file t8_gtest_adapt_partition_balance.cxx
+ * Tests for the adapt, partition and balance routines of mesh handles. 
+ * For the adapt routine, we use the callback and user data of tutorial step 3 as example.
  * The adaptation criterion is to look at the midpoint coordinates of the current element and if
  * they are inside a sphere around a given midpoint we refine, if they are outside, we coarsen. 
+ * The test compares the results of the mesh handle to a forest adapted with the same criterion and balanced and partitioned similarly.
+ * Therefore, the check is based on the assumption that the forest functionality works as intended and is tested elsewhere.
  */
 #include <gtest/gtest.h>
 #include <t8.h>
@@ -98,8 +100,9 @@ forest_adapt_callback_example (t8_forest_t forest, t8_forest_t forest_from, t8_l
   return 0;
 }
 
-/** Test the adapt routine of a mesh handle. 
- * We compare the result to a classically adapted forest with similar callback.
+/** Test the adapt partition and balance routine of a mesh handle. 
+ * The test compares the results of the mesh handle to a forest adapted with the same criterion and balanced and partitioned similarly.
+ * Therefore, the check is based on the assumption that the forest functionality works as intended and is tested elsewhere.
  */
 TEST (t8_gtest_handle_adapt, compare_adapt_with_forest)
 {
@@ -124,11 +127,29 @@ TEST (t8_gtest_handle_adapt, compare_adapt_with_forest)
     mesh_class::mesh_adapt_callback_wrapper<dummy_user_data> (adapt_callback_test<mesh_class>, user_data), false);
   mesh_handle.commit ();
   // Adapt forest classically.
-  forest = t8_forest_new_adapt (forest, forest_adapt_callback_example, 0, 1, &user_data);
+  forest = t8_forest_new_adapt (forest, forest_adapt_callback_example, 0, 0, &user_data);
 
   // Compare results.
   EXPECT_TRUE (t8_forest_is_equal (mesh_handle.get_forest (), forest));
 
+  // Adapt the mesh handle again and apply partition and balance.
+  mesh_handle.set_balance ();
+  mesh_handle.set_partition ();
+  mesh_handle.set_adapt (
+    mesh_class::mesh_adapt_callback_wrapper<dummy_user_data> (adapt_callback_test<mesh_class>, user_data), false);
+  mesh_handle.commit ();
+  EXPECT_TRUE (mesh_handle.is_balanced ());
+
+  // Compare the results again to an appropriate forest.
+  t8_forest_t forest_compare;
+  t8_forest_init (&forest_compare);
+  t8_forest_set_user_data (forest_compare, &user_data);
+  t8_forest_set_adapt (forest_compare, forest, forest_adapt_callback_example, false);
+  t8_forest_set_partition (forest_compare, NULL, false);
+  t8_forest_set_balance (forest_compare, NULL, false);
+  t8_forest_commit (forest_compare);
+  EXPECT_TRUE (t8_forest_is_equal (mesh_handle.get_forest (), forest_compare));
+
   // Clean up.
-  t8_forest_unref (&forest);
+  t8_forest_unref (&forest_compare);
 }