shift_and_combine#

rojak.core.indexing.shift_and_combine(target_array, *, combine_func=<function <lambda>>, shift_dim='pressure_level', offset_start=1, offset_end=1, shift_fill=nan)[source]#

Combine each element along shift_dim with the elements offset_start before and offset_end after it

target_array is shifted forward by offset_start and backward by offset_end along shift_dim, and the two shifted copies are combined element-wise using combine_func. The result is that, at each position, the value at offset_start before it is combined with the value at offset_end after it. The boundary positions which would have required data outside of target_array (and were therefore filled with shift_fill) are trimmed from the result.

With the default combine_func (logical OR), this is useful for e.g. dilating a boolean mask along shift_dim so that a value is also marked True if either of its neighbours (at the given offsets) is True.

Parameters:

  • target_array (T) – Array to shift and combine

  • combine_func (Callable[[T, T], T]) – Function used to combine the two shifted copies of target_array. Defaults to logical OR.

  • shift_dim (str) – Coordinate/dimension to shift along. Defaults to "pressure_level".

  • offset_start (int) – Number of positions to look before each element. Must be non-negative. Defaults to 1.

  • offset_end (int) – Number of positions to look after each element. Must be non-negative. Defaults to 1.

  • shift_fill (Any) – Fill value used at the boundary introduced by shifting, before it is trimmed away. Defaults to np.nan.

Returns:

target_array combined with its shifted copies, with the invalid boundary positions removed

Raises:

ValueError – If offset_start or offset_end is negative, or if shift_dim is not a coordinate of target_array

Return type:

T

>>> import xarray as xr
>>> array = xr.DataArray([False, False, True, False, False], dims="x", coords={"x": [0, 1, 2, 3, 4]})
>>> shift_and_combine(array, shift_dim="x").values
array([ True, False,  True])