dragonfly_sim.utils.mpi_utils

Functions

allgather_numpy(comm, local_arr, global_size)	Allgatherv a rank-local float64 array into a global array on every rank.
assign_global_dof_offset(comm, local_count)	Compute the global DOF offset for this rank via exclusive prefix-sum.

Module Contents

dragonfly_sim.utils.mpi_utils.allgather_numpy(comm, local_arr, global_size)

Allgatherv a rank-local float64 array into a global array on every rank.

Each rank contributes a (possibly different-length) slice of the global array. The slices are ordered by rank: rank 0 occupies [0, len_0), rank 1 occupies [len_0, len_0+len_1), etc. This ordering matches the contiguous DOF blocks produced by assign_global_dof_offset.

Parameters:

commmpi4py.MPI.Comm

local_arrarray-like, float64, 1-D

The local contribution on this rank. Cast to C-contiguous float64 before sending.

global_sizeint

Total number of float64 elements summed across all ranks. Must be the same value on every rank.

Returns:

np.ndarray, float64, shape (global_size,)

The gathered array, identical on every rank.

dragonfly_sim.utils.mpi_utils.assign_global_dof_offset(comm, local_count)

Compute the global DOF offset for this rank via exclusive prefix-sum.

Assigns a contiguous block of global indices to each rank:

rank 0 → [0, local_count_0) rank 1 → [local_count_0, local_count_0 + local_count_1) …

Implementation uses MPI Exscan (exclusive scan with SUM) to compute the offset for each rank in a single collective, then Allreduce on a scalar to obtain the global total. This avoids allocating an N-rank array and is the canonical MPI pattern for assigning contiguous global index ranges to processes.

Parameters:

commmpi4py.MPI.Comm

local_countint

Number of local DOFs on this rank.

Returns:

offsetint

Index of the first global DOF owned by this rank.

global_totalint

Total DOF count across all ranks.