Respect masked array pixels in zonal statistics

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
 

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:

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]))

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"