diff --git a/stackstac/ops.py b/stackstac/ops.py index 778ffd8..924d585 100644 --- a/stackstac/ops.py +++ b/stackstac/ops.py @@ -3,8 +3,7 @@ import xarray as xr -# TODO `fill_value`s besides NaN -def _mosaic(arr, axis, reverse: bool = False): +def _mosaic(arr, axis, reverse: bool = False, nodata: Union[int, float] = np.nan): ax_length = arr.shape[axis] # "normal" means last -> first, "reversed" means first -> last, @@ -14,7 +13,8 @@ def _mosaic(arr, axis, reverse: bool = False): for i in indices: layer = np.take(arr, i, axis=axis) - out = np.where(np.isnan(out), layer, out) + where = np.isnan(out) if np.isnan(nodata) else out == nodata + out = np.where(where, layer, out) return out @@ -23,9 +23,10 @@ def mosaic( dim: Union[None, Hashable, Sequence[Hashable]] = None, axis: Union[None, int, Sequence[int]] = 0, reverse: bool = False, + nodata: Union[int, float] = np.nan, ): """ - Flatten a dimension of a `~xarray.DataArray` by picking the first non-NaN pixel. + Flatten a dimension of a `~xarray.DataArray` by picking the first valid pixel. The order of mosaicing is from last to first, meaning the last item is on top. @@ -41,10 +42,19 @@ def mosaic( reverse: If False (default), the last item along the dimension is on top. If True, the first item in the dimension is on top. + nodata: + The value to treat as invalid. Default: NaN. Returns ------- xarray.DataArray: The mosaicked `~xarray.DataArray`. """ - return arr.reduce(_mosaic, dim=dim, axis=axis, keep_attrs=True, reverse=reverse) + return arr.reduce( + _mosaic, + dim=dim, + axis=axis, + keep_attrs=True, + reverse=reverse, + nodata=nodata, + )