From b9cada56a3ec6d0a5befedcebed617c801f8db2a Mon Sep 17 00:00:00 2001 From: Nicolas Aunai Date: Thu, 24 Oct 2024 10:43:09 +0200 Subject: [PATCH] convenient utils for mpi/hierarchies --- .../pharesee/hierarchy/hierarchy_utils.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/pyphare/pyphare/pharesee/hierarchy/hierarchy_utils.py b/pyphare/pyphare/pharesee/hierarchy/hierarchy_utils.py index 9788f3fb4..95926cd3c 100644 --- a/pyphare/pyphare/pharesee/hierarchy/hierarchy_utils.py +++ b/pyphare/pyphare/pharesee/hierarchy/hierarchy_utils.py @@ -40,6 +40,35 @@ } +def nbr_ranks(hier): + """ + returns the number of mpi ranks used in the given hierarchy + """ + max_rank = 0 + t0 = hier.times()[0] + for _, lvl in hier.levels(t0).items(): + for patch in lvl.patches: + rank = patch.attrs["mpi_rank"] + if rank > max_rank: + max_rank = rank + return max_rank + + +def patch_per_rank(hier): + """ + returns the number of patch per mpi rank for each time step + """ + nbranks = nbr_ranks(hier) + ppr = {} + for t in hier.times(): + ppr[t] = {ir: 0 for ir in np.arange(nbranks + 1)} + for _, lvl in hier.levels(t).items(): + for patch in lvl.patches: + ppr[t][patch.attrs["mpi_rank"]] += 1 + + return ppr + + def are_compatible_hierarchies(hierarchies): ref = hierarchies[0] same_box = True