From 5139c95a03d85cbeef298479e796d7f69473b03e Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Thu, 21 May 2026 19:14:05 -0700 Subject: [PATCH 1/2] rasterize: propagate NaN burns under max/min on every backend (#2255) The GPU paths initialised the max/min output buffer to -inf / +inf and relied on atomicMax / atomicMin to merge writes. NaN comparisons in CUDA return False, so a NaN burn never displaced the inf initialiser (single-NaN case showed -inf instead of NaN) and any finite competitor silently dropped NaN burns (overlap case showed the finite value instead of NaN). The CPU paths had a milder version of the same bug: a NaN burn that arrived after a finite burn was silently dropped because 'NaN > finite' is False, leaving the prior finite pixel in place. Align all four backends on strict NaN propagation: any NaN burn, regardless of input order, poisons the pixel under max / min. CPU checks the burn for NaN explicitly. GPU repurposes the 'written' buffer (already passed through every kernel) as a NaN-mask sized H*W int32 -- atomicMax(mask, 1) records NaN burns and _gpu_finalize_buffers stamps NaN into matching pixels. Updates the asymmetric pins from #2255 deep-sweep pass 2 to assert the unified contract, and adds cross-backend coverage for finite-then-NaN order, single-NaN burns, the min variant, and untouched-pixel fill preservation. --- xrspatial/rasterize.py | 94 ++++++++--- .../test_rasterize_coverage_2026_05_21.py | 149 ++++++++++-------- 2 files changed, 158 insertions(+), 85 deletions(-) diff --git a/xrspatial/rasterize.py b/xrspatial/rasterize.py index e36aaacb..f6da8df3 100644 --- a/xrspatial/rasterize.py +++ b/xrspatial/rasterize.py @@ -90,15 +90,36 @@ def _merge_overwrite(pixel, props, is_first): @ngjit def _merge_max(pixel, props, is_first): - if is_first or props[0] > pixel: - return props[0] + val = props[0] + # NaN burn poisons max regardless of order. Without this branch the + # outcome would depend on whether the NaN burn happened to land before + # or after the finite burns ('NaN > finite' is False, so a NaN burn + # arriving second is silently dropped). Issue #2255. + if val != val: + return val + if is_first: + return val + # A NaN already in the pixel sticks: 'finite > NaN' is False so the + # comparison below would return the pixel anyway, but checking explicitly + # documents the intent and keeps the contract symmetric with the GPU path. + if pixel != pixel: + return pixel + if val > pixel: + return val return pixel @ngjit def _merge_min(pixel, props, is_first): - if is_first or props[0] < pixel: - return props[0] + val = props[0] + if val != val: + return val + if is_first: + return val + if pixel != pixel: + return pixel + if val < pixel: + return val return pixel @@ -1318,10 +1339,25 @@ def _apply_merge_gpu(out, written, order, r, c, props, new_idx): cuda.atomic.add(out, (r, c), 1.0) cuda.atomic.max(order, (r, c), new_idx) elif atomic_mode == 'min': - cuda.atomic.min(out, (r, c), props[0]) + val = props[0] + # CUDA atomicMin uses a < comparison, which returns False for + # any NaN operand, so a NaN burn would silently leave the +inf + # initialiser in place. Track NaN burns in ``written`` (sized + # H*W int32 for min/max, see _gpu_init_buffers) and stamp NaN + # into ``out`` at finalize. Issue #2255. + if val == val: + cuda.atomic.min(out, (r, c), val) + else: + cuda.atomic.max(written, (r, c), 1) cuda.atomic.max(order, (r, c), new_idx) elif atomic_mode == 'max': - cuda.atomic.max(out, (r, c), props[0]) + val = props[0] + # See min branch above; the -inf initialiser would otherwise + # win against any NaN burn. Issue #2255. + if val == val: + cuda.atomic.max(out, (r, c), val) + else: + cuda.atomic.max(written, (r, c), 1) cuda.atomic.max(order, (r, c), new_idx) elif atomic_mode == 'first': # Pass 1: resolve winning input index per pixel. @@ -1882,12 +1918,22 @@ def _gpu_init_buffers(height, width, fill, merge_name): sentinel value used in ``order`` so the caller can blend ``fill`` back in afterwards. - The atomic kernels never consult ``written`` (the legacy - user-callable path is the only consumer), so atomic modes get a - placeholder ``(1, 1)`` buffer to satisfy the kernel signature - without spending ``H*W`` bytes on dead storage. + Most atomic kernels never consult ``written`` (the legacy + user-callable path is the only consumer of the H*W int8 mask), so + they get a placeholder ``(1, 1)`` buffer to satisfy the kernel + signature without spending ``H*W`` bytes on dead storage. + + ``min`` / ``max`` repurpose ``written`` as a NaN-burn mask: the + +/-inf initialisers used by the atomic min/max kernels would + otherwise outrank any NaN burn ('NaN > finite' is False), silently + dropping NaN-valued geometries. The mask records which pixels saw + a NaN burn so finalize can stamp NaN back in. Issue #2255. + ``cuda.atomic.max`` requires int32 / int64 / uint32 / uint64 (int8 + is not supported), hence the dtype bump for these two modes. """ - if merge_name in _GPU_ATOMIC_MERGES: + if merge_name in ('min', 'max'): + written = cupy.zeros((height, width), dtype=cupy.int32) + elif merge_name in _GPU_ATOMIC_MERGES: written = cupy.zeros((1, 1), dtype=cupy.int8) else: written = cupy.zeros((height, width), dtype=cupy.int8) @@ -1924,11 +1970,16 @@ def _gpu_init_buffers(height, width, fill, merge_name): return out, written, order, order_sentinel -def _gpu_finalize_buffers(out, order, order_sentinel, fill, merge_name): +def _gpu_finalize_buffers(out, written, order, order_sentinel, fill, + merge_name): """Apply the post-pass blend that restores ``fill`` for untouched pixels (sum/count/min/max/first/last) and replaces the +/- inf initialisers used by min/max. + For ``min`` / ``max``, also stamps NaN into pixels that received at + least one NaN burn -- ``written`` is repurposed as the NaN-mask for + those modes (see ``_gpu_init_buffers``). Issue #2255. + Returns ``out`` (modified in-place where convenient). """ if merge_name is None or order_sentinel is None: @@ -1937,10 +1988,13 @@ def _gpu_finalize_buffers(out, order, order_sentinel, fill, merge_name): if merge_name in ('sum', 'count'): # Pixels with no contribution should show ``fill``, not 0. out = cupy.where(touched, out, cupy.float64(fill)) - elif merge_name == 'min': - # +inf initialiser leaks into untouched pixels. - out = cupy.where(touched, out, cupy.float64(fill)) - elif merge_name == 'max': + elif merge_name in ('min', 'max'): + # Any pixel that saw a NaN burn must show NaN, otherwise the + # +/-inf initialiser would leak through (or the finite max/min + # of the non-NaN burns would silently win). Issue #2255. + nan_burned = written > 0 + out = cupy.where(nan_burned, cupy.float64(cupy.nan), out) + # +/-inf initialiser leaks into untouched pixels; restore fill. out = cupy.where(touched, out, cupy.float64(fill)) # first/last already initialise ``out`` to ``fill``; pass-2 only # writes for touched pixels, so no blend is needed. @@ -2069,8 +2123,8 @@ def _launch_phase(k_scan, k_lines, k_boundary, k_points): kernels['burn_lines_supercover_pass2'], kernels['burn_points_pass2']) - final_out = _gpu_finalize_buffers(out, order, order_sentinel, fill, - atomic_mode) + final_out = _gpu_finalize_buffers(out, written, order, order_sentinel, + fill, atomic_mode) return final_out.astype(dtype) @@ -2531,8 +2585,8 @@ def _launch_phase(k_scan, k_lines, k_boundary, k_points): kernels['burn_lines_supercover_pass2'], kernels['burn_points_pass2']) - final_out = _gpu_finalize_buffers(out, order, order_sentinel, fill, - atomic_mode) + final_out = _gpu_finalize_buffers(out, written, order, order_sentinel, + fill, atomic_mode) return final_out.astype(dtype) diff --git a/xrspatial/tests/test_rasterize_coverage_2026_05_21.py b/xrspatial/tests/test_rasterize_coverage_2026_05_21.py index 6a5bcce9..b6f825c5 100644 --- a/xrspatial/tests/test_rasterize_coverage_2026_05_21.py +++ b/xrspatial/tests/test_rasterize_coverage_2026_05_21.py @@ -239,14 +239,10 @@ class TestNaNBurnValues: leave the fill value in covered cells, which is a different observable than emitting NaN there. - NOTE: the GPU ``max`` / ``min`` merges currently suppress NaN burn - values, asymmetric with the CPU IEEE-propagating behaviour pinned - here. See issue #2255. The - ``test_nan_burn_overlaps_max_gpu_suppresses_nan`` and - ``test_nan_burn_single_geom_max_gpu_returns_neg_inf`` tests below - pin the current (asymmetric) GPU observable so the divergence is - visible in CI; both will need to be inverted once the GPU kernels - are aligned with CPU semantics. + Issue #2255 aligned the cross-backend contract on strict NaN + propagation for ``max`` / ``min``: any NaN burn poisons the output + pixel regardless of order. See ``TestNaNPropagationAcrossBackends`` + below for the cross-backend pins. """ @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) @@ -296,80 +292,103 @@ def test_nan_burn_line(self, backend_name, kw): f"got per-row NaN counts {nan_rows}" ) - # CPU backends (numpy / dask+numpy) follow IEEE: max(NaN, 1.0) returns - # NaN because the comparison ``1.0 > NaN`` is False, so the kernel - # keeps the prior NaN pixel. GPU backends (cupy / dask+cupy) currently - # return 1.0: the GPU kernel initialises the output buffer to -inf for - # max merge and uses atomic_max which is NaN-suppressing under IEEE - # device semantics. See issue #2255. The CPU pin and the GPU - # asymmetry pin keep the divergence visible until the GPU kernels are - # aligned with the CPU is_first semantics. - @pytest.mark.parametrize('backend_name,kw', [ - ALL_BACKENDS[0], # numpy - ALL_BACKENDS[2], # dask_numpy - ]) - def test_nan_burn_overlaps_max_cpu_propagates(self, backend_name, kw): - """CPU: ``max(NaN, 1.0) == NaN``. +# --------------------------------------------------------------------------- +# Issue #2255: strict NaN propagation under max / min on every backend +# --------------------------------------------------------------------------- - Pin the IEEE NaN-propagating behaviour so a future change that - switches the CPU max to ``fmax`` (NaN-suppressing) is visible as - a diff, not silent. - """ + +class TestNaNPropagationMaxMin: + """Cross-backend pin: any NaN burn poisons ``max`` / ``min``. + + Before the fix the GPU paths silently dropped NaN burns + (``atomicMax`` / ``atomicMin`` compare with ``<`` / ``>``, both + False for NaN), and the CPU path dropped a NaN burn that arrived + *after* a finite burn (``NaN > finite`` is False). Issue #2255 + aligned all four backends on strict NaN propagation: any NaN burn, + regardless of input order, makes the pixel NaN. + """ + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_nan_first_then_finite_max(self, backend_name, kw): r = rasterize( [(box(0, 0, 10, 5), np.nan), (box(0, 0, 10, 5), 1.0)], width=10, height=5, bounds=(0, 0, 10, 5), fill=0, merge='max', **kw) data = _materialise(r) - assert np.all(np.isnan(data)) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected all-NaN, got {data}") - @pytest.mark.parametrize('backend_name,kw', [ - ALL_BACKENDS[1], # cupy - ALL_BACKENDS[3], # dask_cupy - ]) - def test_nan_burn_overlaps_max_gpu_suppresses_nan(self, backend_name, kw): - """GPU: ``max(NaN, 1.0) == 1.0`` (NaN-suppressing). - - Asymmetric with the CPU behaviour above; see issue #2255. - Pinning the current observable here so the GPU<->CPU divergence - is documented in CI until the GPU max kernel can be aligned with - IEEE NaN propagation. - """ + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_finite_then_nan_max(self, backend_name, kw): + """CPU previously kept the finite value here ("NaN > finite" is + False so the order-dependent kernel dropped the later NaN). + After the fix the NaN burn explicitly poisons the pixel.""" r = rasterize( - [(box(0, 0, 10, 5), np.nan), (box(0, 0, 10, 5), 1.0)], + [(box(0, 0, 10, 5), 1.0), (box(0, 0, 10, 5), np.nan)], width=10, height=5, bounds=(0, 0, 10, 5), fill=0, merge='max', **kw) data = _materialise(r) - # Current (asymmetric) GPU behaviour: NaN was suppressed and the - # finite 1.0 won. This pin will need to flip to ``np.all(isnan)`` - # once the GPU max kernel is fixed to match the CPU IEEE - # propagation. - np.testing.assert_array_equal(data, np.ones_like(data)) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected all-NaN, got {data}") - @pytest.mark.parametrize('backend_name,kw', [ - ALL_BACKENDS[1], # cupy - ALL_BACKENDS[3], # dask_cupy - ]) - def test_nan_burn_single_geom_max_gpu_returns_neg_inf( - self, backend_name, kw): - """GPU: a single NaN-burn polygon under ``max`` leaves the - kernel's -inf init value in covered pixels. - - On CPU the value would be NaN (the burn was the only write and - ``_merge_max`` returns ``props[0]`` on first-write). On GPU the - buffer is initialised to -inf and atomicMax(-inf, NaN) keeps -inf - because NaN comparisons are False. See issue #2255. Pin so the - divergence is visible against the CPU behaviour pinned by - test_nan_burn_polygon[numpy]. - """ + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_single_nan_burn_max(self, backend_name, kw): + """Lone NaN burn lands as NaN, not -inf (GPU init) and not the + fill (CPU default).""" r = rasterize([(box(0, 0, 10, 5), np.nan)], width=10, height=5, bounds=(0, 0, 10, 5), fill=0, merge='max', **kw) data = _materialise(r) - # The 5x10 covered region is -inf, not NaN, not 0. - assert np.all(np.isneginf(data)), ( - f"{backend_name}: expected -inf in covered pixels, got {data}" - ) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected NaN, got {data}") + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_nan_first_then_finite_min(self, backend_name, kw): + """Same poisoning contract under ``merge='min'``.""" + r = rasterize( + [(box(0, 0, 10, 5), np.nan), (box(0, 0, 10, 5), 1.0)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='min', **kw) + data = _materialise(r) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected all-NaN, got {data}") + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_single_nan_burn_min(self, backend_name, kw): + """Single NaN burn under ``merge='min'`` -> NaN (was +inf on GPU).""" + r = rasterize([(box(0, 0, 10, 5), np.nan)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='min', **kw) + data = _materialise(r) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected NaN, got {data}") + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_partial_nan_partial_finite_max(self, backend_name, kw): + """Non-overlapping polygons: NaN poisons only its own pixels, + finite polygon's pixels keep their finite max. Guards against a + regression where the NaN-mask leaks across the raster.""" + r = rasterize( + [(box(0, 0, 5, 5), np.nan), (box(5, 0, 10, 5), 1.0)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='max', **kw) + data = _materialise(r) + assert np.all(np.isnan(data[:, :5])), ( + f"{backend_name}: left half should be NaN") + np.testing.assert_array_equal(data[:, 5:], np.ones((5, 5))) + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_untouched_pixels_get_fill_under_max(self, backend_name, kw): + """Pixels with no burn at all retain the fill, even when other + pixels saw NaN burns. Guards against the NaN-mask bleeding into + untouched cells during finalize.""" + r = rasterize([(box(0, 0, 5, 5), np.nan)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=-9.0, merge='max', **kw) + data = _materialise(r) + assert np.all(np.isnan(data[:, :5])) + np.testing.assert_array_equal(data[:, 5:], np.full((5, 5), -9.0)) # --------------------------------------------------------------------------- From 54fb3818144ae2b10dc8c1d56454317ef10f9338 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Thu, 21 May 2026 19:18:59 -0700 Subject: [PATCH 2/2] Address review nits: drop dead NaN branch, move tests to dedicated file (#2255) - Removed the redundant 'if pixel != pixel: return pixel' branch in _merge_max / _merge_min. The earlier 'val != val' early-return rules out NaN burns reaching the check, and a NaN pixel from an earlier burn is already preserved by the 'val > pixel' (or 'val < pixel') comparison returning False against NaN, falling through to 'return pixel'. A short comment documents the implicit stickiness. - Moved TestNaNPropagationMaxMin out of the deep-sweep coverage file (whose docstring claims 'No source changes') into a dedicated test_rasterize_nan_propagation_2255.py file, since the contract it pins required a source-side fix. --- xrspatial/rasterize.py | 9 +- .../test_rasterize_coverage_2026_05_21.py | 98 ----------- .../test_rasterize_nan_propagation_2255.py | 161 ++++++++++++++++++ 3 files changed, 163 insertions(+), 105 deletions(-) create mode 100644 xrspatial/tests/test_rasterize_nan_propagation_2255.py diff --git a/xrspatial/rasterize.py b/xrspatial/rasterize.py index f6da8df3..bffeeff7 100644 --- a/xrspatial/rasterize.py +++ b/xrspatial/rasterize.py @@ -99,11 +99,8 @@ def _merge_max(pixel, props, is_first): return val if is_first: return val - # A NaN already in the pixel sticks: 'finite > NaN' is False so the - # comparison below would return the pixel anyway, but checking explicitly - # documents the intent and keeps the contract symmetric with the GPU path. - if pixel != pixel: - return pixel + # If pixel is already NaN from an earlier burn, 'val > NaN' is False + # below and we fall through to 'return pixel', keeping NaN sticky. if val > pixel: return val return pixel @@ -116,8 +113,6 @@ def _merge_min(pixel, props, is_first): return val if is_first: return val - if pixel != pixel: - return pixel if val < pixel: return val return pixel diff --git a/xrspatial/tests/test_rasterize_coverage_2026_05_21.py b/xrspatial/tests/test_rasterize_coverage_2026_05_21.py index b6f825c5..df063d9d 100644 --- a/xrspatial/tests/test_rasterize_coverage_2026_05_21.py +++ b/xrspatial/tests/test_rasterize_coverage_2026_05_21.py @@ -293,104 +293,6 @@ def test_nan_burn_line(self, backend_name, kw): ) -# --------------------------------------------------------------------------- -# Issue #2255: strict NaN propagation under max / min on every backend -# --------------------------------------------------------------------------- - - -class TestNaNPropagationMaxMin: - """Cross-backend pin: any NaN burn poisons ``max`` / ``min``. - - Before the fix the GPU paths silently dropped NaN burns - (``atomicMax`` / ``atomicMin`` compare with ``<`` / ``>``, both - False for NaN), and the CPU path dropped a NaN burn that arrived - *after* a finite burn (``NaN > finite`` is False). Issue #2255 - aligned all four backends on strict NaN propagation: any NaN burn, - regardless of input order, makes the pixel NaN. - """ - - @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) - def test_nan_first_then_finite_max(self, backend_name, kw): - r = rasterize( - [(box(0, 0, 10, 5), np.nan), (box(0, 0, 10, 5), 1.0)], - width=10, height=5, bounds=(0, 0, 10, 5), - fill=0, merge='max', **kw) - data = _materialise(r) - assert np.all(np.isnan(data)), ( - f"{backend_name}: expected all-NaN, got {data}") - - @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) - def test_finite_then_nan_max(self, backend_name, kw): - """CPU previously kept the finite value here ("NaN > finite" is - False so the order-dependent kernel dropped the later NaN). - After the fix the NaN burn explicitly poisons the pixel.""" - r = rasterize( - [(box(0, 0, 10, 5), 1.0), (box(0, 0, 10, 5), np.nan)], - width=10, height=5, bounds=(0, 0, 10, 5), - fill=0, merge='max', **kw) - data = _materialise(r) - assert np.all(np.isnan(data)), ( - f"{backend_name}: expected all-NaN, got {data}") - - @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) - def test_single_nan_burn_max(self, backend_name, kw): - """Lone NaN burn lands as NaN, not -inf (GPU init) and not the - fill (CPU default).""" - r = rasterize([(box(0, 0, 10, 5), np.nan)], - width=10, height=5, bounds=(0, 0, 10, 5), - fill=0, merge='max', **kw) - data = _materialise(r) - assert np.all(np.isnan(data)), ( - f"{backend_name}: expected NaN, got {data}") - - @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) - def test_nan_first_then_finite_min(self, backend_name, kw): - """Same poisoning contract under ``merge='min'``.""" - r = rasterize( - [(box(0, 0, 10, 5), np.nan), (box(0, 0, 10, 5), 1.0)], - width=10, height=5, bounds=(0, 0, 10, 5), - fill=0, merge='min', **kw) - data = _materialise(r) - assert np.all(np.isnan(data)), ( - f"{backend_name}: expected all-NaN, got {data}") - - @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) - def test_single_nan_burn_min(self, backend_name, kw): - """Single NaN burn under ``merge='min'`` -> NaN (was +inf on GPU).""" - r = rasterize([(box(0, 0, 10, 5), np.nan)], - width=10, height=5, bounds=(0, 0, 10, 5), - fill=0, merge='min', **kw) - data = _materialise(r) - assert np.all(np.isnan(data)), ( - f"{backend_name}: expected NaN, got {data}") - - @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) - def test_partial_nan_partial_finite_max(self, backend_name, kw): - """Non-overlapping polygons: NaN poisons only its own pixels, - finite polygon's pixels keep their finite max. Guards against a - regression where the NaN-mask leaks across the raster.""" - r = rasterize( - [(box(0, 0, 5, 5), np.nan), (box(5, 0, 10, 5), 1.0)], - width=10, height=5, bounds=(0, 0, 10, 5), - fill=0, merge='max', **kw) - data = _materialise(r) - assert np.all(np.isnan(data[:, :5])), ( - f"{backend_name}: left half should be NaN") - np.testing.assert_array_equal(data[:, 5:], np.ones((5, 5))) - - @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) - def test_untouched_pixels_get_fill_under_max(self, backend_name, kw): - """Pixels with no burn at all retain the fill, even when other - pixels saw NaN burns. Guards against the NaN-mask bleeding into - untouched cells during finalize.""" - r = rasterize([(box(0, 0, 5, 5), np.nan)], - width=10, height=5, bounds=(0, 0, 10, 5), - fill=-9.0, merge='max', **kw) - data = _materialise(r) - assert np.all(np.isnan(data[:, :5])) - np.testing.assert_array_equal(data[:, 5:], np.full((5, 5), -9.0)) - - # --------------------------------------------------------------------------- # Cat 1 MEDIUM -- Nested GeometryCollection # --------------------------------------------------------------------------- diff --git a/xrspatial/tests/test_rasterize_nan_propagation_2255.py b/xrspatial/tests/test_rasterize_nan_propagation_2255.py new file mode 100644 index 00000000..0886dd1d --- /dev/null +++ b/xrspatial/tests/test_rasterize_nan_propagation_2255.py @@ -0,0 +1,161 @@ +"""Cross-backend NaN-propagation pins for ``rasterize`` ``merge='max'`` / ``'min'``. + +Issue #2255: before the fix the GPU paths silently dropped NaN burns +(``atomicMax`` / ``atomicMin`` compare with ``<`` / ``>``, both False +for NaN operands), and the CPU path dropped a NaN burn that arrived +*after* a finite burn (``NaN > finite`` is False, so the kernel kept +the prior finite pixel). The PR aligned all four backends on strict +NaN propagation: any NaN burn, regardless of input order, makes the +pixel NaN under max / min. + +These tests pin that contract across numpy / cupy / dask+numpy / +dask+cupy so a future regression (e.g. swapping ``fmax`` / ``fmin`` +back in, or restoring the -inf-init optimisation without the NaN-mask +finalize) fails loudly. +""" +from __future__ import annotations + +import numpy as np +import pytest + +try: + from shapely.geometry import box + has_shapely = True +except ImportError: + has_shapely = False + +if has_shapely: + from xrspatial.rasterize import rasterize + +pytestmark = pytest.mark.skipif( + not has_shapely, reason="shapely not installed" +) + +try: + import cupy + has_cupy = True +except ImportError: + has_cupy = False + +try: + import dask # noqa: F401 + has_dask = True +except ImportError: + has_dask = False + +try: + from numba import cuda + has_cuda = has_cupy and cuda.is_available() +except Exception: + has_cuda = False + +skip_no_cuda = pytest.mark.skipif( + not has_cuda, reason="CUDA / CuPy not available") +skip_no_dask = pytest.mark.skipif( + not has_dask, reason="dask not installed") + + +def _materialise(result): + """Compute (if dask) and copy to host (if cupy) so callers see ndarray.""" + data = result.data + if hasattr(data, 'compute'): + data = data.compute() + if has_cupy and isinstance(data, cupy.ndarray): + return cupy.asnumpy(data) + return np.asarray(data) + + +ALL_BACKENDS = [ + pytest.param('numpy', {}, id='numpy'), + pytest.param('cupy', {'use_cuda': True}, + marks=skip_no_cuda, id='cupy'), + pytest.param('dask_numpy', {'chunks': (5, 5)}, + marks=skip_no_dask, id='dask_numpy'), + pytest.param('dask_cupy', {'use_cuda': True, 'chunks': (5, 5)}, + marks=[skip_no_cuda, skip_no_dask], id='dask_cupy'), +] + + +class TestNaNPropagationMaxMin: + """Any NaN burn poisons ``max`` / ``min``, regardless of order.""" + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_nan_first_then_finite_max(self, backend_name, kw): + r = rasterize( + [(box(0, 0, 10, 5), np.nan), (box(0, 0, 10, 5), 1.0)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='max', **kw) + data = _materialise(r) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected all-NaN, got {data}") + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_finite_then_nan_max(self, backend_name, kw): + """CPU previously kept the finite value here ('NaN > finite' is + False so the order-dependent kernel dropped the later NaN). + After the fix the NaN burn explicitly poisons the pixel.""" + r = rasterize( + [(box(0, 0, 10, 5), 1.0), (box(0, 0, 10, 5), np.nan)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='max', **kw) + data = _materialise(r) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected all-NaN, got {data}") + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_single_nan_burn_max(self, backend_name, kw): + """Lone NaN burn lands as NaN, not -inf (GPU init) and not the + fill (CPU default).""" + r = rasterize([(box(0, 0, 10, 5), np.nan)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='max', **kw) + data = _materialise(r) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected NaN, got {data}") + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_nan_first_then_finite_min(self, backend_name, kw): + """Same poisoning contract under ``merge='min'``.""" + r = rasterize( + [(box(0, 0, 10, 5), np.nan), (box(0, 0, 10, 5), 1.0)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='min', **kw) + data = _materialise(r) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected all-NaN, got {data}") + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_single_nan_burn_min(self, backend_name, kw): + """Single NaN burn under ``merge='min'`` -> NaN (was +inf on GPU).""" + r = rasterize([(box(0, 0, 10, 5), np.nan)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='min', **kw) + data = _materialise(r) + assert np.all(np.isnan(data)), ( + f"{backend_name}: expected NaN, got {data}") + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_partial_nan_partial_finite_max(self, backend_name, kw): + """Non-overlapping polygons: NaN poisons only its own pixels, + finite polygon's pixels keep their finite max. Guards against a + regression where the NaN-mask leaks across the raster.""" + r = rasterize( + [(box(0, 0, 5, 5), np.nan), (box(5, 0, 10, 5), 1.0)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=0, merge='max', **kw) + data = _materialise(r) + assert np.all(np.isnan(data[:, :5])), ( + f"{backend_name}: left half should be NaN") + np.testing.assert_array_equal(data[:, 5:], np.ones((5, 5))) + + @pytest.mark.parametrize('backend_name,kw', ALL_BACKENDS) + def test_untouched_pixels_get_fill_under_max(self, backend_name, kw): + """Pixels with no burn retain the fill, even when other pixels + saw NaN burns. Guards against the NaN-mask bleeding into + untouched cells during finalize.""" + r = rasterize([(box(0, 0, 5, 5), np.nan)], + width=10, height=5, bounds=(0, 0, 10, 5), + fill=-9.0, merge='max', **kw) + data = _materialise(r) + assert np.all(np.isnan(data[:, :5])) + np.testing.assert_array_equal(data[:, 5:], np.full((5, 5), -9.0))