From e3312701c289279bbdbc1b1e6e9a71a040c8c2ba Mon Sep 17 00:00:00 2001
From: Lars Pastewka <lars.pastewka@imtek.uni-freiburg.de>
Date: Wed, 21 Oct 2026 23:49:09 +0200
Subject: [PATCH] ENH: Gather timing summaries through a muGrid communicator
 directly

A communicator could gather the per-rank timers only if it was mpi4py's
or wrapped one (muGrid's mpi4py_comm). muGrid 1.6 gathers Python objects
itself -- gather(obj, root=0) with mpi4py's semantics -- and no longer
needs mpi4py, so a communicator whose gather takes root is used as it is.
muGrid's earlier gather, an array concatenation without root, is still
left alone.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 CHANGELOG.md        |  5 +++++
 muTimer/Timer.py    | 38 ++++++++++++++++++++++++++------------
 tests/test_timer.py | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 65c26c1..3d89aeb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+### Changed
+- Timing summaries are gathered across ranks through a muGrid communicator
+  (1.6 or newer) directly; mpi4py is no longer needed for that
+
 ## [1.1.0] - 2026-06-11
 ### Added
 - MPI awareness
diff --git a/muTimer/Timer.py b/muTimer/Timer.py
index 554287e..6e9d03c 100644
--- a/muTimer/Timer.py
+++ b/muTimer/Timer.py
@@ -1,5 +1,6 @@
 # MIT License. See LICENSE file for details.
 
+import inspect
 import json
 import time
 from contextlib import contextmanager
@@ -19,19 +20,32 @@ def _comm_rank_size(comm):
     return comm.rank, comm.size  # muGrid-style properties
 
 
-def _find_object_gather_comm(comm):
-    """Find a communicator capable of gathering Python objects.
+def _gathers_objects(comm):
+    """Whether ``comm.gather(obj, root=0)`` gathers Python objects.
 
-    Returns an mpi4py-style communicator (one with a lowercase ``gather``
-    method that pickles arbitrary objects), or None if the given
-    communicator provides no such interface.
+    True for mpi4py communicators and for muGrid communicators from 1.6 on,
+    whose ``gather`` takes ``root``; muGrid's earlier ``gather`` was an
+    array concatenation without ``root`` and is left alone.
     """
-    if hasattr(comm, "gather") and hasattr(comm, "Get_rank"):
+    gather = getattr(comm, "gather", None)
+    if gather is None:
+        return False
+    if hasattr(comm, "Get_rank"):  # mpi4py's interface
+        return True
+    try:
+        return "root" in inspect.signature(gather).parameters
+    except (TypeError, ValueError):  # no introspectable signature
+        return False
+
+
+def _find_object_gather_comm(comm):
+    """Find a communicator capable of gathering Python objects: the given
+    one, or the mpi4py communicator a muGrid communicator was made from
+    (``mpi4py_comm``); None if there is neither."""
+    if _gathers_objects(comm):
         return comm
-    # muGrid's CommunicatorWrapper exposes the underlying mpi4py
-    # communicator as ``mpi4py_comm``
     candidate = getattr(comm, "mpi4py_comm", None)
-    if hasattr(candidate, "gather") and hasattr(candidate, "Get_rank"):
+    if candidate is not None and _gathers_objects(candidate):
         return candidate
     return None
 
@@ -376,9 +390,9 @@ def print_summary(self, title=None, name_width=30, comm=None):
         is available (either passed here or to the constructor), timing data
         is gathered to rank 0 and a single summary is printed there, showing
         the mean, minimum and maximum time across ranks; all other ranks
-        print nothing. Gathering requires an mpi4py communicator (muGrid
-        communicators created from one work too); without it, rank 0 prints
-        its local timings only.
+        print nothing. Gathering needs a communicator whose ``gather`` takes
+        Python objects -- mpi4py's, or muGrid's from 1.6 on; without one,
+        rank 0 prints its local timings only.
 
         Args:
             title: Title for the summary section (default: "Timing Summary")
diff --git a/tests/test_timer.py b/tests/test_timer.py
index 607bc95..5024ccb 100644
--- a/tests/test_timer.py
+++ b/tests/test_timer.py
@@ -301,3 +301,36 @@ def test_serial_mpi4py_style_comm_keeps_serial_format(capsys):
     out = capsys.readouterr().out
     assert "Total" in out and "Average" in out
     assert "MPI processes" not in out
+
+
+def test_object_gather_detection():
+    from muTimer.Timer import _find_object_gather_comm
+
+    class Mpi4pyLike:
+        def Get_rank(self):
+            return 0
+
+        def gather(self, obj, root=0):
+            return [obj]
+
+    class MuGridLike:  # muGrid >= 1.6: gather takes root, no Get_rank
+        rank = 0
+
+        def gather(self, obj, root=0):
+            return [obj]
+
+    class OldMuGridLike:  # muGrid < 1.6: gather concatenated arrays
+        rank = 0
+        mpi4py_comm = None
+
+        def gather(self, arg):
+            return arg
+
+    class Wrapped(OldMuGridLike):
+        mpi4py_comm = Mpi4pyLike()
+
+    assert _find_object_gather_comm(Mpi4pyLike()) is not None
+    assert _find_object_gather_comm(MuGridLike()) is not None
+    assert _find_object_gather_comm(OldMuGridLike()) is None
+    assert isinstance(_find_object_gather_comm(Wrapped()), Mpi4pyLike)
+    assert _find_object_gather_comm(object()) is None