diff --git a/xvec/accessor.py b/xvec/accessor.py index d63edeb..2aacc70 100644 --- a/xvec/accessor.py +++ b/xvec/accessor.py @@ -658,7 +658,8 @@ def mask( xr.DataArray A DataArray with the same shape as the original, where the elements matching the predicate are set to True. """ - cube_data = self._obj.data.ravel() + # need to replace nan with None + cube_data = self._obj.where(~self._obj.isnull(), None).data.ravel() tree = shapely.STRtree(cube_data) indices = tree.query(geometry, predicate=predicate, distance=distance) if indices.ndim == 1: diff --git a/xvec/zonal.py b/xvec/zonal.py index 2a7b384..32eed1a 100644 --- a/xvec/zonal.py +++ b/xvec/zonal.py @@ -581,17 +581,22 @@ def _get_mean( ): from rasterio import features - mask = features.geometry_mask( - [geom_arr.item()], - out_shape=( - obj[y_coords].shape[0], - obj[x_coords].shape[0], - ), - transform=transform, - invert=True, - all_touched=all_touched, - ) - masked = obj.where(xr.DataArray(mask, dims=(y_coords, x_coords))) + if pd.isna(geom_arr.item()): + masked = obj.where( + xr.DataArray(np.full_like(obj.data, False), dims=(y_coords, x_coords)) + ) + else: + mask = features.geometry_mask( + [geom_arr.item()], + out_shape=( + obj[y_coords].shape[0], + obj[x_coords].shape[0], + ), + transform=transform, + invert=True, + all_touched=all_touched, + ) + masked = obj.where(xr.DataArray(mask, dims=(y_coords, x_coords))) if nodata is not None: masked = masked.where(masked != nodata)