pygetm.parallel module

class pygetm.parallel.DistributedArray(tiling: Tiling, field: ndarray, halox: int, haloy: int, overlap: int = 0)

Bases: object

compare_halos(group: Neighbor = Neighbor.ALL) → bool

group2task: list[tuple[list[Prequest], list[Prequest], list[tuple[ndarray, ndarray]], list[tuple[ndarray, ndarray]]]]

halo2name

rank

update_halos(group: Neighbor = Neighbor.ALL)

update_halos_finish(group: Neighbor = Neighbor.ALL)

update_halos_start(group: Neighbor = Neighbor.ALL)

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.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

• 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 objetc 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

wrap(*args, **kwargs) → DistributedArray

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, max_aspect_ratio: int = 2, weight_unmasked: int = 2, weight_any: int = 1, weight_halo: int = 10, max_protrude: float = 0.5, 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)

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()