pyuoi.mpi_utils

Helper functions for loading data and managing arrays across ranks with MPI.

pyuoi.mpi_utils.Bcast_from_root(send, comm=None, root=0)[source]

Broadcast an array from root to all MPI ranks.

Parameters

  • send (ndarray or None) – Array to send from root to all ranks. send in other ranks has no effect.

  • comm (MPI.COMM_WORLD) – MPI communicator.

  • root (int) – This rank contains the array to send.

Returns

send – Each rank will have a copy of the array from root.

Return type

ndarray

pyuoi.mpi_utils.Gatherv_rows(send, comm=None, root=0)[source]

Concatenate arrays along the first axis using Gatherv on root.

Parameters

  • send (ndarray) – The arrays to concatenate. All dimensions must be equal except for the first.

  • comm (MPI.COMM_WORLD) – MPI communicator.

  • root (int) – This rank will contain the Gatherv’ed array.

Returns

rec – Gatherv’ed array on root or None on other ranks.

Return type

ndarray or None

pyuoi.mpi_utils.check_valid_ndarray(X)[source]

Checks whether X is a ndarray and returns a contiguous version.

pyuoi.mpi_utils.load_data_MPI(h5_name, X_key='X', y_key='y', comm=None, root=0)[source]

Load data from an HDF5 file and broadcast it across MPI ranks.

This is a helper function. It is also possible to load the data without this function.

Parameters

  • h5_name (str) – Path to h5 file.

  • X_key (str) – Key for the features dataset. (default: ‘X’)

  • y_key (str) – Key for the targets dataset. (default: ‘y’)

  • comm (MPI.COMM_WORLD) – MPI communicator.

  • root (int) – This rank will load the data from file.

Returns

  • X (ndarray) – Features on all MPI ranks.

  • y (ndarray) – Targets on all MPI ranks.