Rename option to view_patch_sizes, include the average patch size

firedrakeproject · Jan 14, 2025 · 84c9409 · 84c9409
1 parent 8699c3c
commit 84c9409
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/firedrake/preconditioners/asm.py b/firedrake/preconditioners/asm.py
@@ -92,13 +92,17 @@ def initialize(self, pc):
         self.asmpc = asmpc
 
         self._patch_statistics = []
-        if opts.getBool("print_patch_statistics", default=False):
+        if opts.getBool("view_patch_sizes", default=False):
             # Compute and stash patch statistics
+            mpi_comm = pc.comm.tompi4py()
             max_local_patch = max(is_.getSize() for is_ in ises)
             min_local_patch = min(is_.getSize() for is_ in ises)
-            max_global_patch = pc.comm.tompi4py().allreduce(max_local_patch, op=MPI.MAX)
-            min_global_patch = pc.comm.tompi4py().allreduce(min_local_patch, op=MPI.MIN)
-            msg = f"Minimum / maximum patch sizes: {min_global_patch} / {max_global_patch}\n"
+            sum_local_patch = sum(is_.getSize() for is_ in ises)
+            max_global_patch = mpi_comm.allreduce(max_local_patch, op=MPI.MAX)
+            min_global_patch = mpi_comm.allreduce(min_local_patch, op=MPI.MIN)
+            sum_global_patch = mpi_comm.allreduce(sum_local_patch, op=MPI.SUM)
+            avg_global_patch = sum_global_patch / mpi_comm.allreduce(len(ises) if sum_local_patch > 0 else 0, op=MPI.SUM)
+            msg = f"Minimum / average / maximum patch sizes : {min_global_patch} / {avg_global_patch} / {max_global_patch}\n"
             self._patch_statistics.append(msg)
 
     @abc.abstractmethod

diff --git a/tests/firedrake/regression/test_star_pc.py b/tests/firedrake/regression/test_star_pc.py
@@ -51,8 +51,8 @@ def test_star_equivalence(problem_type, backend):
                        "mg_levels_ksp_max_it": 1,
                        "mg_levels_pc_type": "python",
                        "mg_levels_pc_python_type": "firedrake.ASMStarPC",
-                       "mg_levels_pc_star_print_patch_statistics": None,
-                       "mg_levels_pc_star_construct_dim": 0}
+                       "mg_levels_pc_star_construct_dim": 0,
+                       "mg_levels_pc_star_view_patch_sizes": None}
 
         comp_params = {"mat_type": "aij",
                        "snes_type": "ksponly",