qttools.utils.mpi_utils#
source module qttools.utils.mpi_utils
Functions
-
get_section_sizes — Computes the number of un-evenly divided elements per section.
-
distributed_load — Loads an array from disk and distributes it to all ranks.
-
get_local_slice — Returns the local slice of a distributed array.
source get_section_sizes(num_elements: int, num_sections: int = comm.size, strategy: str = 'balanced') → tuple[list, int]
Computes the number of un-evenly divided elements per section.
Parameters
-
num_elements : int — The total number of elements to divide.
-
num_sections : int, optional — The number of sections to divide the elements into. Defaults to the number of MPI ranks.
-
strategy : str, optional — The strategy to use for dividing the elements. Can be one of "balanced" (default) or "greedy". In the "balanced" strategy, the elements are divided as evenly as possible across the sections. In the "greedy" strategy, the elements are divided such that the we get many sections with the maximum number of elements.
Returns
-
section_sizes : list — The sizes of each section.
-
effective_num_elements : int — The effective number of elements after sectioning.
Examples
Raises
-
ValueError
source distributed_load(path: Path) → sparse.spmatrix | NDArray
Loads an array from disk and distributes it to all ranks.
Parameters
-
path : Path — The path to the file to load.
Returns
-
sparse.spmatrix | NDArray — The loaded array.
Raises
-
FileNotFoundError — Occurs on every rank where the file does not exist.
-
ValueError
source get_local_slice(global_array: NDArray, comm: MPI.Comm = comm) → NDArray
Returns the local slice of a distributed array.
Parameters
-
global_array : NDArray — The global array to slice.
Returns
-
NDArray — The local slice of the global array.