stems.parallel module

Helper tools to facilitate turning computation into dask-able gufuncs

References

[1] https://docs.scipy.org/doc/numpy/reference/c-api.generalized-ufuncs.html [2] http://numba.pydata.org/numba-doc/dev/user/vectorize.html

stems.parallel.iter_chunks(dim_sizes, chunk_sizes=None)[source]

Return slices that help iterate over dimensions in chunks/blocks

Parameters

  • dim_sizes (dict[str, int]) – Dimension sizes

  • chunk_sizes (None, int, or dict[str, int]) – Sizes of chunks. If int, will use the same size for all dimensions. If None, will iter over each element in all dimensions (e.g., one pixel)

Yields

dict[str, slice or int] – Mapping of dimension name to select statement (e.g., {'x': slice(0, 50), 'y': slice(0, 50)})

stems.parallel.iter_noncore_chunks(data, core_dims, chunk_sizes=None)[source]

Yield data in chunks selected out from non-core dimensions

Parameters

  • data (xr.DataArray or xr.Dataset) – Data

  • core_dims (str or list[str]) – One or more dimensions to exclude from selection. These dimensions will be the only dimensions on the data returned

  • chunk_size (int, dict[str, int], or None) – Chunk/block/window size when iterating over noncore dimensions. If None, will return a slice for each element in each dimension.

Yields

dict[str, slice] – isel select statements to retrieve data

stems.parallel.map_collect_1d(core_dims, arg_idx=0, concat_axis=0, concat_func=<function concatenate at 0x7f760da74840>)[source]

Decorator that maps a function across pixels and concatenates the result

Parameters

  • core_dims (Sequence[str]) – Core dimensions to exclude from broadcast/selection. For example, if applying a time series model over all x/y in a DataArray, the core_dim could be ‘time’.

  • arg_idx (int or tuple[int], optional) – Index(es) of *args to pick as the data to iterate over.

  • concat_axis (int, optional) – Axis along which the arrays will be joined (passed to concat_func)

  • concat_func (callable, optional) – Function used to concatenate the data. Should take the axis keyword

Returns

Decorator for function (parametrized by core_dims and arg_idx)

Return type

callable