Drop the per-cell NaN early-out in the planar CPU slope kernel (#3742)

* Drop the per-cell NaN early-out in the planar CPU slope kernel (#3739)

_cpu started every interior cell with `if np.isnan(data[y, x]): continue`.
On DEMs with scattered nodata that branch is data-dependent and
mispredicts more than the few float32 ops and one arctan it skips.

Remove the branch and fold the centre back in with a select after the
arithmetic. Neighbour NaN already propagates through the Horn stencil,
so the select only has to cover a NaN centre with valid neighbours.
Results are bitwise identical to the old kernel on 2000x4000 float32
DEMs with no NaN, 30% random NaN and a NaN left half, and on 3x3, 1xN
and int16 inputs. On the 30% random NaN raster the kernel goes from
~92 ms to ~55 ms; the NaN-free and contiguous-NaN cases are within
noise.

Add a SlopeNaN asv benchmark with 30% random NaN over the interior so
the nodata path is timed from now on. get_xr_dataarray(include_nan=True)
only sets the [0, 0] corner to NaN, which the kernel never visits, so
the benchmark adds its own speckle.

Add tests pinning the NaN footprint on a speckled raster: NaN at every
centre-NaN cell and its 8-neighbours, finite elsewhere, on numpy and
dask+numpy.

* Address review notes on the slope NaN early-out PR (#3739)

Note in SlopeNaN that its dask timing forces the compute while
Slope.time_slope does not, so the two classes' dask numbers are not
comparable. Build the speckle mask from a float32 draw to halve the
setup allocation at nx=10000. State in the footprint test comment that
the tests pass on the old kernel too, since they pin behaviour rather
than reproduce a regression. Drop the stale "(unchanged)" from the
planar section banner.

Author: brendancol

benchmarks/benchmarks/slope.py

@@ -1,6 +1,8 @@
+import numpy as np
+
 from xrspatial import slope
 
-from .common import Benchmarking
+from .common import Benchmarking, get_xr_dataarray
 
 
 class Slope(Benchmarking):
@@ -9,3 +11,31 @@ def __init__(self):
 
     def time_slope(self, nx, type):
         return self.time(nx, type)
+
+
+class SlopeNaN(Benchmarking):
+    # Speckled nodata: the planar CPU kernel's cost depends on whether the
+    # NaN pattern is predictable, so time it separately from Slope.
+    # get_xr_dataarray(include_nan=True) only sets the [0, 0] corner to NaN,
+    # which is on the border the kernel never visits, so add 30% random NaN
+    # over the interior on top of it.
+    params = ([100, 300, 1000, 3000, 10000], ["numpy", "dask"])
+    param_names = ("nx", "type")
+
+    def __init__(self):
+        super().__init__(func=slope)
+
+    def setup(self, nx, type):
+        ny = nx // 2
+        agg = get_xr_dataarray((ny, nx), type, include_nan=True)
+        rng = np.random.default_rng(71942)
+        speckle = rng.random((ny, nx), dtype=np.float32) < 0.3
+        self.xr = agg.where(~speckle)
+
+    def time_slope_nan(self, nx, type):
+        # Force the compute so the dask case times the kernel rather than
+        # graph construction. Slope.time_slope does not, so the dask numbers
+        # of the two classes are not directly comparable.
+        result = self.func(self.xr)
+        if hasattr(result.data, "compute"):
+            result.data.compute()

xrspatial/slope.py

@@ -43,7 +43,7 @@ def _geodesic_cuda_dims(shape):
 
 
 # =====================================================================
-# Planar backend functions (unchanged)
+# Planar backend functions
 # =====================================================================
 
 @ngjit
@@ -54,8 +54,11 @@ def _cpu(data, cellsize_x, cellsize_y):
     rows, cols = data.shape
     for y in range(1, rows - 1):
         for x in range(1, cols - 1):
-            if np.isnan(data[y, x]):
-                continue
+            # No early-out on a NaN centre: the branch is unpredictable on
+            # speckled nodata and costs more than the arithmetic it skips.
+            # Neighbour NaN already propagates through the stencil; the centre
+            # is folded back in with a select so out[y, x] is NaN whenever
+            # data[y, x] is.
             a = data[y + 1, x - 1]
             b = data[y + 1, x]
             c = data[y + 1, x + 1]
@@ -67,7 +70,9 @@ def _cpu(data, cellsize_x, cellsize_y):
             dz_dx = ((c + 2 * f + i) - (a + 2 * d + g)) / (8 * cellsize_x)
             dz_dy = ((g + 2 * h + i) - (a + 2 * b + c)) / (8 * cellsize_y)
             p = (dz_dx * dz_dx + dz_dy * dz_dy) ** .5
-            out[y, x] = np.arctan(p) * 57.29578
+            r = np.arctan(p) * 57.29578
+            ctr = data[y, x]
+            out[y, x] = r if ctr == ctr else np.nan
     return out
 
 

xrspatial/tests/test_slope.py

@@ -625,3 +625,49 @@ def test_units_overridden_to_degrees_geodesic():
     result = slope(raster, method='geodesic')
     assert result.attrs['units'] == 'degrees'
     assert raster.attrs['units'] == 'm'
+
+
+# The planar CPU kernel has no per-cell NaN early-out (#3739): neighbour NaN
+# flows through the Horn stencil and the centre is masked with a select. Pin
+# the resulting footprint on speckled nodata so a future edit can't widen or
+# shrink it: NaN at every centre-NaN cell and its 8-neighbours, finite
+# everywhere else in the interior. The old kernel produced the same footprint,
+# so these pass on either version; they pin behaviour rather than reproduce a
+# regression.
+def _speckled_nan_data():
+    rng = np.random.default_rng(3739)
+    data = rng.normal(0, 2, size=(40, 60))
+    data[rng.random(data.shape) < 0.1] = np.nan
+    return data
+
+
+def _expected_nan_footprint(data):
+    nan_mask = np.isnan(data)
+    footprint = np.zeros_like(nan_mask)
+    for dy in (-1, 0, 1):
+        for dx in (-1, 0, 1):
+            footprint[1:-1, 1:-1] |= nan_mask[1 + dy:data.shape[0] - 1 + dy,
+                                              1 + dx:data.shape[1] - 1 + dx]
+    footprint[0, :] = footprint[-1, :] = footprint[:, 0] = footprint[:, -1] = True
+    return footprint
+
+
+def test_speckled_nan_footprint_numpy():
+    data = _speckled_nan_data()
+    agg = create_test_raster(data, backend='numpy', attrs={'res': (1, 1)})
+    result = slope(agg)
+    general_output_checks(agg, result, verify_attrs=False)
+    np.testing.assert_array_equal(np.isnan(result.data), _expected_nan_footprint(data))
+    assert np.all(np.isfinite(result.data[~np.isnan(result.data)]))
+
+
+@dask_array_available
+def test_speckled_nan_footprint_dask_numpy():
+    data = _speckled_nan_data()
+    numpy_agg = create_test_raster(data, backend='numpy', attrs={'res': (1, 1)})
+    dask_agg = create_test_raster(data, backend='dask+numpy',
+                                  attrs={'res': (1, 1)}, chunks=(16, 24))
+    assert_numpy_equals_dask_numpy(numpy_agg, dask_agg, slope, nan_edges=False,
+                                   verify_attrs=False)
+    np.testing.assert_array_equal(np.isnan(slope(dask_agg).data.compute()),
+                                  _expected_nan_footprint(data))