Python bindings: Add Dataset.ReadAsMaskedArray

Author: dbaston

autotest/gcore/numpy_rw.py

@@ -983,13 +983,15 @@ def test_numpy_rw_band_read_as_array_getlasterrormsg():
 def test_numpy_rw_masked_array_1():
 
     ds = gdal.Open("data/byte.tif")
-
     band = ds.GetRasterBand(1)
 
     masked_arr = band.ReadAsMaskedArray()
-
     assert not numpy.any(masked_arr.mask)
 
+    ds_masked_arr = ds.ReadAsMaskedArray()
+    assert not numpy.any(ds_masked_arr.mask)
+    assert numpy.all(masked_arr == ds_masked_arr)
+
 
 def test_numpy_rw_masked_array_2():
 
@@ -1007,6 +1009,90 @@ def test_numpy_rw_masked_array_2():
 
     assert masked_arr.sum() == arr[mask == 255].sum()
 
+    ds_masked_arr = ds.ReadAsMaskedArray()
+    assert numpy.all(ds_masked_arr.mask == masked_arr.mask)
+    assert numpy.all(ds_masked_arr.data == masked_arr.data)
+
+
+###############################################################################
+# Test dataset read into a masked array
+
+
+def test_numpy_rw_masked_array_3():
+
+    nx = 6
+    ny = 4
+    nz = 3
+    nodata = -999
+
+    ds = gdal.GetDriverByName("MEM").Create("", nx, ny, nz, eType=gdal.GDT_Int16)
+    for i in range(nz):
+        ds.GetRasterBand(i + 1).SetNoDataValue(nodata)
+
+    data = numpy.arange(nx * ny * nz).reshape(nz, ny, nx).astype(numpy.int16)
+    data[:, 1, 1] = nodata
+    ds.WriteArray(data)
+
+    masked_data = numpy.ma.masked_array(data, mask=data == nodata)
+
+    assert numpy.ma.allequal(masked_data, ds.ReadAsMaskedArray())
+
+    # read a single band
+    assert numpy.ma.allequal(masked_data[1, :, :], ds.ReadAsMaskedArray(band_list=[2]))
+
+    # read multiple bands
+    assert numpy.ma.allequal(
+        numpy.ma.stack((masked_data[2, :, :], masked_data[0, :, :])),
+        ds.ReadAsMaskedArray(band_list=[3, 1]),
+    )
+
+    # read a sub-window (square)
+    assert numpy.ma.allequal(
+        numpy.ma.stack((masked_data[2, 0:2, 1:3], masked_data[0, 0:2, 1:3])),
+        ds.ReadAsMaskedArray(band_list=[3, 1], xsize=2, ysize=2, xoff=1, yoff=0),
+    )
+
+    # read a sub-window (single pixel)
+    assert numpy.ma.allequal(
+        numpy.ma.stack((masked_data[2, 1:2, 1:2], masked_data[0, 1:2, 1:2])),
+        ds.ReadAsMaskedArray(band_list=[3, 1], xsize=1, ysize=1, xoff=1, yoff=1),
+    )
+
+    # read a sub-window (read 2x2 square into 4x4 square)
+    assert numpy.ma.allequal(
+        numpy.ma.repeat(
+            numpy.ma.repeat(
+                numpy.ma.stack((masked_data[2, 0:2, 1:3], masked_data[0, 0:2, 1:3])),
+                2,
+                axis=2,
+            ),
+            2,
+            axis=1,
+        ),
+        ds.ReadAsMaskedArray(
+            band_list=[3, 1], xsize=2, ysize=2, xoff=1, yoff=0, buf_xsize=4, buf_ysize=4
+        ),
+    )
+
+    # read a sub-window (read 2x2 square into single pixel)
+    assert numpy.ma.allequal(
+        numpy.ma.mean(
+            numpy.ma.stack((masked_data[2, 0:2, 1:3], masked_data[0, 0:2, 1:3])),
+            axis=(1, 2),
+            keepdims=True,
+        ).round(),
+        ds.ReadAsMaskedArray(
+            band_list=[3, 1],
+            xsize=2,
+            ysize=2,
+            xoff=1,
+            yoff=0,
+            buf_xsize=1,
+            buf_ysize=1,
+            resample_alg=gdal.GRIORA_Average,
+        ),
+    )
+
 
 ###############################################################################
 # Test type code mapping

swig/include/python/gdal_python.i

@@ -1110,6 +1110,42 @@ CPLErr ReadRaster1( double xoff, double yoff, double xsize, double ysize,
 
 %pythoncode %{
 
+    def ReadAsMaskedArray(self, xoff=0, yoff=0, xsize=None, ysize=None, 
+                    buf_xsize=None, buf_ysize=None, buf_type=None,
+                    resample_alg=gdalconst.GRIORA_NearestNeighbour,
+                    callback=None,
+                    callback_data=None,
+                    band_list=None):
+        """
+        Read a window from raster bands into a NumPy masked array.
+
+        Parameters are the same as for :py:meth:`ReadAsArray`.
+        """
+
+        import numpy as np
+
+        arr = self.ReadAsArray(xoff=xoff, yoff=yoff, xsize=xsize, ysize=ysize, 
+                               buf_xsize=buf_xsize, buf_ysize=buf_ysize, buf_type=buf_type,
+                               resample_alg=resample_alg, band_list=band_list)
+        
+        if band_list is None:
+            band_list = [i+1 for i in range(self.RasterCount)]
+        
+        all_valid = all(self.GetRasterBand(band).GetMaskFlags() == GMF_ALL_VALID for band in band_list)
+
+        if all_valid:
+            return np.ma.masked_array(arr, False)
+       
+        masks = [self.GetRasterBand(band).GetMaskBand().ReadAsArray(
+                xoff=xoff, yoff=yoff,
+                win_xsize=xsize, win_ysize=ysize,
+                buf_xsize=buf_xsize, buf_ysize=buf_ysize,
+                resample_alg=gdalconst.GRIORA_Mode) != 255
+                 for band in band_list]
+
+        return np.ma.masked_array(arr, np.vstack(masks))
+
+
     def ReadAsArray(self, xoff=0, yoff=0, xsize=None, ysize=None, buf_obj=None,
                     buf_xsize=None, buf_ysize=None, buf_type=None,
                     resample_alg=gdalconst.GRIORA_NearestNeighbour,