Allow bigger partition discrepancies, include missing import

illinois-ceesd · Aug 22, 2024 · dd8f535 · dd8f535
1 parent 21a9681
commit dd8f535
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/bin/meshdist.py b/bin/meshdist.py
@@ -195,7 +195,8 @@ def get_mesh_data():
     def my_partitioner(mesh, tag_to_elements, num_ranks):
         from mirgecom.simutil import geometric_mesh_partitioner
         return geometric_mesh_partitioner(
-            mesh, num_ranks, auto_balance=True, debug=False)
+            mesh, num_ranks, auto_balance=True, debug=True,
+            imbalance_tolerance=.2)
 
     part_func = my_partitioner if use_1d_part else None
 

diff --git a/mirgecom/simutil.py b/mirgecom/simutil.py
@@ -94,7 +94,10 @@
     DiscretizationDOFAxisTag
 )
 from arraycontext import flatten, map_array_container
-from grudge.discretization import DiscretizationCollection
+from grudge.discretization import (
+    DiscretizationCollection,
+    PartID
+)
 from grudge.dof_desc import DD_VOLUME_ALL
 from meshmode.dof_array import DOFArray
 from collections import defaultdict