MPI (taurex.mpi)

Module for wrapping MPI functions.

Most functions will do nothing if mpi4py is not present and will simply replicate expected behaviour. This allows for MPI and non-MPI code to be run without any changes (In theory…).

class Ops(value)[source]

Bases: str, Enum

Available MPI operations

MAX = 'max'
MIN = 'min'
PROD = 'prod'
SUM = 'sum'
allgather(value: T) → List[T][source]

Gathers all values from all processes or returns value if mpi is not installed

Parameters:

value – Value to gather

Returns:

List of gathered values or list containing value if mpi is not installed

Return type:

list

allocate_as_shared(arr: ndarray, logger: Logger | None = None, force_shared: bool | None = False)[source]

Converts a numpy array into an MPI shared memory.

This allow for things like opacities to be loaded only once per node when using MPI. Only activates if mpi4py installed and when enabled via the mpi_use_shared input:

[Global]
mpi_use_shared = True

or force_shared=True otherwise does nothing and returns the same array back

Parameters:

  • arr (numpy array) – Array to convert

  • logger (Logger) – Logger object to print outputs

  • force_shared (bool) – Force conversion to shared memory

Returns:

If enabled and MPI present, shared memory version of array otherwise the original array

Return type:

array

allreduce(value: T, op: Ops) → T[source]

Reduces all values from all processes or returns value if mpi is not installed

Parameters:

  • value – Value to reduce

  • op (Ops) – Operation to perform

Returns:

Reduced value or value if mpi is not installed

Return type:

result

barrier(comm: Any | None = None) → None[source]

Waits for all processes to finish.

Does nothing if mpi4py not present

Parameters:

comm (int, optional) – MPI communicator, default is MPI_COMM_WORLD

broadcast(array: T, rank: int | None = 0) → T[source]

Broadcasts array from rank or returns array if mpi is not installed

Parameters:

  • array – Array to broadcast

  • rank (int, optional) – Rank to broadcast from, default is 0

Returns:

Broadcasted array or array if mpi is not installed

Return type:

array

convert_op(operation: Ops) → Any[source]

Converts string to MPI operation.

get_rank(comm: Any = None) → int[source]

Gets rank or returns 0 if mpi is not installed

Parameters:

comm (int, optional) – MPI communicator, default is MPI_COMM_WORLD

Returns:

Rank of process in communitor or 0 if MPI is not installed

Return type:

int

has_mpi() → bool[source]

Checks if mpi4py is installed

Returns:

True if mpi4py is installed, False otherwise

Return type:

bool

nprocs() → int[source]

Gets number of processes or returns 1 if mpi is not installed

Returns:

Rank of process or 1 if MPI is not installed

Return type:

int

only_master_rank(f) → Callable[source]

A decorator to ensure only the master MPI rank can run it.

shared_comm() → Any[source]

Return shared memory communicator.

Returns the process id within a node. Used for shared memory.

shared_rank() → int[source]

Gets rank within shared memory communicator. (MPI only)