pygetm.parallel module

class pygetm.parallel.BaseHaloUpdater

Bases: object

finish()

start()

class pygetm.parallel.Gather(tiling: Tiling, shape: tuple[int, ...], dtype: DTypeLike, fill_value=None, root: int = 0, overlap: int = 0)

Bases: object

class pygetm.parallel.GatherFromIndices(comm: Comm, local2global: ndarray, index_shape: tuple[int, ...], nonindexed_shape: tuple[int, ...], dtype: DTypeLike, *, fill_value=None, trailing_index: bool = True)

Bases: object

Gather values at specific indices from all ranks to rank 0.

The target indices can have any shape (0D for a station, 1D for a transect). Any remaining (non-indexed) dimensions can be present too. The values contained in the local subdomain are specified by local2global.

Parameters:

• comm – MPI communicator

• local2global – 1D array specifying the location of points from the local subdomain in the global index. This is 1D array of indices into the flattened target index array. Values therefore must lie between 0 and index_shape.prod().

• index_shape – shape of the target indices in the global domain

• nonindexed_shape – shape of any non-indexed dimensions

• dtype – data type of the values to be gathered

• fill_value – value to use for indexed points that are not present in any subdomain

• trailing_index – whether the indexed dimensions come last (default). If False, they are assumed to come first.

__call__(locvalues: ndarray, globvalues: ndarray | None = None, globslice=()) → ndarray

Gather values from each subdomain.

class pygetm.parallel.HaloUpdater(rank: int)

Bases: BaseHaloUpdater

Class to manage halo exchange for a single variable. It creates the data buffers and persistent requests for MPI communication, and it provides methods to start and finish the halo exchange.

__call__()

Perform halo exchange by starting all communication, waiting for it to complete, and updating the target arrays with the received data. This is equivalent to calling start() followed by finish().

add_recv(neighbor: str, req: Prequest, outer: ndarray, cache: ndarray)

add_send(neighbor: str, req: Prequest, inner: ndarray, cache: ndarray)

all_reqs: list[Prequest]

compare() → bool

Verify that values in halos are in sync with the corresponding inner values of the neighboring subdomains. This can be used for debugging purposes.

finish()

Finish halo exchange by waiting for all communication to complete and updating the target arrays with the received data. This must be called after start()

freeze()

halo2name: dict[int, str]

rank

recv_data: list[tuple[ndarray, ndarray]]

recv_reqs: list[Prequest]

send_data: list[tuple[ndarray, ndarray]]

send_reqs: list[Prequest]

start()

Start halo exchange. This can be used to overlap communication with computation. finish() must be called after the computation to ensure that the halo exchange is completed and the target array is updated with the received data.

class pygetm.parallel.Neighbor(*values)

Bases: IntEnum

ALL = 0

BOTTOM = 2

BOTTOMLEFT = 1

BOTTOMRIGHT = 3

LEFT = 4

LEFT_AND_RIGHT = 10

LEFT_AND_RIGHT_AND_TOP_AND_BOTTOM = 12

RIGHT = 5

TOP = 7

TOPLEFT = 6

TOPRIGHT = 8

TOP_AND_BOTTOM = 9

TOP_AND_RIGHT = 11

class pygetm.parallel.Scatter(tiling: Tiling, field: ndarray, halox: int, haloy: int, share: int = 0, fill_value=0.0, scale: int = 1, exclude_halos: bool = False, root: int = 0)

Bases: object

class pygetm.parallel.Tiling(nrow: int | None = None, ncol: int | None = None, map: ndarray | None = None, comm: Comm | None = None, periodic_x: bool = False, periodic_y: bool = False, ncpus: int | None = None)

Bases: object

__bool__() → bool

Return True if the curent subdomain has any neighbors, False otherwise.

allreduce(field: ~numpy.ndarray, op=<mpi4py.MPI.Op object>) → ndarray

static autodetect(mask: ArrayLike, *, ncpus: int | None = None, logger: Logger | None = None, comm: Comm | None = None, max_protrude: float = 0.5, **kwargs) → Tiling

Auto-detect the optimal subdomain division and return it as Tiling object

Parameters:

• mask – mask for the global domain (0: masked, non-zero: active)

• ncpus – number of cores to use (default: size of MPI communicator)

• logger – logger to use to write diagnostic messages

• comm – MPI communicator to use to parallize the search for the optimal subdomain division (default: MPI.COMM_WORLD). It will also be used to detect the number of cores to use if ncpus is not provided.

• max_protrude – maximum fraction (0-1) of a subdomain that can protrude out of the global domain

• **kwargs – additional keyword arguments to pass to Tiling

describe()

Print neighbours of the current subdomain

dump(path: str)

Save all information about the subdomain division to a pickle file from which a Tiling object can be recreated with Tiling.load()

Parameters:

path – file to save information to

static load(path: str) → Tiling

Create a Tiling object from information in a pickle file that was produced with Tiling.dump()

property nactive: int

plot(ax=None, background: ndarray | None = None, cmap='viridis')

Plot the current subdomain division

Parameters:

• ax – matplotlib.axes.Axes to plot into. If not provided, a new matplotlib.figure.Figure with single axes will be created.

• background – field to use as background for the subdomain division. If provided, it must have the extent of the global domain. If not provided a single filled rectabgle covering the global domain will be used as background.

reduce(field: ~numpy.ndarray, op=<mpi4py.MPI.Op object>) → ndarray | None

report(logger: Logger)

Write information about the subdomain decompostion to the log. Log messages are suppressed if the decompositon only has one subdomain.

set_extent(nx_glob: int, ny_glob: int, nx_sub: int | None = None, ny_sub: int | None = None, xoffset_global: int = 0, yoffset_global: int = 0)

Set extent of the global domain and the subdomain, and the offsets of the very first subdomain.

Parameters:

• nx_glob – x extent of the global domain

• ny_glob – y extent of the global domain

• nx_sub – x extent of a subdomain If not set, infer from number of ranks in single row

• ny_sub – y extent of a subdomain If not set, infer from number of ranks in single column

• xoffset_global – x offset of left-most subdomains

• yoffset_global – y offset of bottom-most subdomains

subdomain2rawslices(irow: int | None = None, icol: int | None = None, *, halox_sub: int = 0, haloy_sub: int = 0, scale: int = 1, share: int = 0) → tuple[slice, slice] | tuple[None, None]

subdomain2slices(irow: int | None = None, icol: int | None = None, *, halox_sub: int = 0, haloy_sub: int = 0, halox_glob: int = 0, haloy_glob: int = 0, scale: int = 1, share: int = 0, exclude_halos: bool = True, exclude_global_halos: bool = False) → tuple[tuple[Any, slice, slice], tuple[Any, slice, slice], tuple[int, int], tuple[int, int]]

Determine the activate extent of a subdomain in terms of the slices spanned in subdomain arrays and in global arrays

Parameters:

• irow – row of the subdomain (default: current)

• icol – column of the subdomain (default: current)

• halo_sub – width of the halos in the subdomain array to be sliced

• halo_glob – width of the halos in the global array to be sliced

• scale – the extent of subdomain array relative to the subdomain extent provided to set_extent() (for instance, 2 for supergrid arrays)

• share – the number of points shared by subdomains (for instance, 1 for supergrid arrays and X grids)

• exclude_halos – whether to excludes the halos of the subdomain from the slice

• exclude_global_halos – whether to excludes the halos of the global domain from the slice

Returns:

Tuple with the local slices, global slices, extent of a subdomain array, and extent of a global array. The extents always include halos and shared points

pygetm.parallel.create_halo_updaters(tiling: Tiling | None, field: ndarray, halox: int, haloy: int, overlap: int = 0) → Sequence[BaseHaloUpdater]

pygetm.parallel.find_all_optimal_divisons(mask: ArrayLike, max_ncpus: int | Iterable[int], **kwargs) → Mapping[str, Any] | None

pygetm.parallel.find_optimal_divison(mask: ArrayLike, *, ncpus: int | None = None, min_cpus: int | None = None, max_aspect_ratio: int = 2, weight_unmasked: int = 2, weight_any: int = 1, weight_halo: int = 10, max_protrude: float = 0.25, logger: Logger | None = None, comm: Comm | None = None) → Mapping[str, Any] | None

pygetm.parallel.get_logger(level=20, comm=<mpi4py.MPI.Intracomm object>) → Logger

pygetm.parallel.mpi4py_autofree(obj: T) → T

pygetm.parallel.test_scaling(setup_script: str, nmax: int | None = None, nmin: int = 1, extra_args: Iterable[str] = [], plot: bool = True, out: str | None = None, compare: Iterable[str] | None = None)

pygetm.parallel.test_scaling_command()