diff --git a/src/rasterstats/io.py b/src/rasterstats/io.py index a07d698..1c62b69 100644 --- a/src/rasterstats/io.py +++ b/src/rasterstats/io.py @@ -292,6 +292,8 @@ def boundless_array(arr, window, nodata, masked=False): # create an array of nodata values out = np.empty(shape=window_shape, dtype=arr.dtype) out[:] = nodata + if np.ma.isMaskedArray(arr): + out = np.ma.array(out, mask=np.zeros(window_shape, dtype=bool)) # Fill with data where overlapping nr_start = olr_start - wr_start @@ -308,7 +310,7 @@ def boundless_array(arr, window, nodata, masked=False): ] if masked: - out = np.ma.MaskedArray(out, mask=(out == nodata)) + out = np.ma.masked_equal(out, nodata, copy=False) return out diff --git a/src/rasterstats/main.py b/src/rasterstats/main.py index 0db7918..68a1a82 100644 --- a/src/rasterstats/main.py +++ b/src/rasterstats/main.py @@ -195,7 +195,9 @@ def gen_zonal_stats( band = band_num with Raster(raster, affine, nodata, band) as rast: - features_iter = read_features(vectors, layer, engine=engine, chunk_size=chunk_size) + features_iter = read_features( + vectors, layer, engine=engine, chunk_size=chunk_size + ) for _, feat in enumerate(features_iter): geom = feat["geometry"] if not isinstance(geom, Geometry): @@ -223,7 +225,10 @@ def gen_zonal_stats( # Mask the source data array # mask everything that is not a valid value or not within our geom - masked = np.ma.MaskedArray(fsrc.array, mask=(isnodata | ~rv_array)) + source_mask = np.ma.getmaskarray(fsrc.array) + masked = np.ma.MaskedArray( + fsrc.array, mask=(source_mask | isnodata | ~rv_array) + ) # If we're on 64 bit platform and the array is an integer type # make sure we cast to 64 bit to avoid overflow for certain numpy ops @@ -310,13 +315,19 @@ def gen_zonal_stats( feature_stats[pctile] = float(np.percentile(pctarr, q)) if "nodata" in stats or "nan" in stats: - featmasked = np.ma.MaskedArray(fsrc.array, mask=(~rv_array)) + featmasked = np.ma.MaskedArray( + fsrc.array, mask=(source_mask | ~rv_array) + ) if "nodata" in stats: - feature_stats["nodata"] = float((featmasked == fsrc.nodata).sum()) + feature_stats["nodata"] = float( + np.ma.filled(featmasked == fsrc.nodata, False).sum() + ) if "nan" in stats: feature_stats["nan"] = ( - float(np.isnan(featmasked).sum()) if has_nan else 0 + float(np.ma.filled(np.isnan(featmasked), False).sum()) + if has_nan + else 0 ) if add_stats is not None: diff --git a/tests/test_io.py b/tests/test_io.py index e2327aa..7ae6914 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -240,6 +240,35 @@ def test_boundless_masked(): assert c.mask.any() and not c.mask.all() +@pytest.mark.parametrize("dim3", [False, True]) +@pytest.mark.parametrize("masked", [False, True]) +def test_boundless_preserves_source_mask(dim3, masked): + source = np.ma.array([[1, 2], [3, 4]], mask=[[False, True], [False, False]]) + if dim3: + source = source[np.newaxis, :, :] + + result = boundless_array( + source, window=((-1, 2), (0, 3)), nodata=-999, masked=masked + ) + + assert isinstance(result, np.ma.MaskedArray) + actual_mask = np.ma.getmaskarray(result) + if dim3: + actual_mask = actual_mask[0] + assert actual_mask[1, 1] + assert result.count() == (3 if masked else 8) + assert actual_mask[0, 0] == masked + assert actual_mask[2, 2] == masked + + +def test_boundless_masked_array_without_overlap(): + source = np.ma.array([[1]], mask=[[True]]) + result = boundless_array(source, window=((2, 3), (2, 3)), nodata=-999, masked=True) + + assert result.shape == (1, 1) + assert result.mask.all() + + def test_window_bounds(): with rasterio.open(raster) as src: win = ((0, src.shape[0]), (0, src.shape[1])) diff --git a/tests/test_zonal.py b/tests/test_zonal.py index d39ed59..4679fad 100644 --- a/tests/test_zonal.py +++ b/tests/test_zonal.py @@ -452,6 +452,45 @@ def test_nan_nodata(): assert stats[0]["max"] == 12.75 +def test_masked_ndarray_excludes_source_mask(): + polygon = Polygon([[0, 0], [2, 0], [2, 2], [0, 2]]) + arr = np.ma.array([[1, 100], [-999, 3]], mask=[[False, True], [False, False]]) + affine = Affine(1, 0, 0, 0, -1, 2) + + stats = zonal_stats( + polygon, arr, affine=affine, nodata=-999, stats="count sum mean nodata" + )[0] + + assert stats == {"count": 2, "sum": 4.0, "mean": 2.0, "nodata": 1.0} + + +def test_masked_ndarray_excludes_masked_nodata_and_nan(): + polygon = Polygon([[0, 0], [2, 0], [2, 2], [0, 2]]) + arr = np.ma.array( + [[1.0, np.nan], [-999.0, 3.0]], + mask=[[False, True], [True, False]], + ) + affine = Affine(1, 0, 0, 0, -1, 2) + + stats = zonal_stats( + polygon, arr, affine=affine, nodata=-999, stats="count sum nodata nan" + )[0] + + assert stats == {"count": 2, "sum": 4.0, "nodata": 0.0, "nan": 0} + + +def test_masked_ndarray_with_every_pixel_masked(): + polygon = Polygon([[0, 0], [2, 0], [2, 2], [0, 2]]) + arr = np.ma.array([[1, 2], [3, 4]], mask=True) + affine = Affine(1, 0, 0, 0, -1, 2) + + stats = zonal_stats( + polygon, arr, affine=affine, nodata=-999, stats="count mean nodata" + )[0] + + assert stats == {"count": 0, "mean": None, "nodata": 0.0} + + def test_some_nodata_ndarray(): polygons = data_dir / "polygons.shp" raster = data_dir / "slope_nodata.tif"