From d5e8533676e766c9b1ccd919cdb096af274b25fb Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Fri, 4 Sep 2026 11:50:09 -0400 Subject: [PATCH 1/2] Replace the D8 code-to-offset if/elif chain with a lookup table (#3738) _code_to_offset mapped a D8 direction code to its (dy, dx) neighbour offset with an eight-way if/elif chain. _flow_accum_cpu calls it twice per cell and stream_link_d8, flow_length_d8, flow_path_d8, stream_order_d8 and hand_d8 import it for their own per-cell loops. Index two module-level int64 arrays by the code instead. numba freezes module-level arrays as compile-time constants, so the call compiles to a range check and two loads. The range check is written as an inside-the-box test (0 <= c <= 128) on purpose: int(nan) is INT64_MIN in numba and there is no bounds checking, so a rejection-form guard would let a NaN code index out of bounds (same class as the erode fix in #3703). Non-power-of-two codes in range and anything outside 0..128 keep returning (0, 0), and the returns stay integer. On a 1000x2000 float64 flow-direction raster _flow_accum_cpu dropped from 56.3 ms to 39.7 ms (median of 7 after warmup) with identical results, on both the clean raster and one built from a DEM with 30% NaN. Add a parametrised test pinning the mapping for all eight codes, 0, in-range non-codes, out-of-range ints, floats and NaN, plus a check that only the eight D8 entries of the tables are populated. --- xrspatial/hydro/flow_accumulation_d8.py | 36 +++++++------ .../hydro/tests/test_flow_accumulation_d8.py | 53 +++++++++++++++++++ 2 files changed, 73 insertions(+), 16 deletions(-) diff --git a/xrspatial/hydro/flow_accumulation_d8.py b/xrspatial/hydro/flow_accumulation_d8.py index 677ad3ef9..e6e8950d9 100644 --- a/xrspatial/hydro/flow_accumulation_d8.py +++ b/xrspatial/hydro/flow_accumulation_d8.py @@ -192,26 +192,30 @@ def _no_weight_cupy(): # Direction helpers # ===================================================================== +# (dy, dx) row/col offset per D8 code, indexed by the code itself. Only +# the eight power-of-two entries are non-zero; every other in-range index +# stays (0, 0) so codes like 3 or 5 keep meaning "no flow". numba freezes +# module-level arrays as compile-time constants, so the lookup below is a +# range check plus two loads instead of an eight-way branch chain. +_D8_DY = np.zeros(129, dtype=np.int64) +_D8_DX = np.zeros(129, dtype=np.int64) +for _code, (_dy, _dx) in ((1, (0, 1)), (2, (1, 1)), (4, (1, 0)), (8, (1, -1)), + (16, (0, -1)), (32, (-1, -1)), (64, (-1, 0)), + (128, (-1, 1))): + _D8_DY[_code] = _dy + _D8_DX[_code] = _dx +del _code, _dy, _dx + + @ngjit def _code_to_offset(code): """Return (dy, dx) row/col offset for a D8 direction code.""" c = int(code) - if c == 1: - return 0, 1 - elif c == 2: - return 1, 1 - elif c == 4: - return 1, 0 - elif c == 8: - return 1, -1 - elif c == 16: - return 0, -1 - elif c == 32: - return -1, -1 - elif c == 64: - return -1, 0 - elif c == 128: - return -1, 1 + # Inside-the-box guard, not a rejection test: int(nan) is INT64_MIN in + # numba and there is no bounds checking, so ``c < 0 or c > 128`` style + # guards would not stop a NaN code from indexing out of bounds. + if 0 <= c <= 128: + return _D8_DY[c], _D8_DX[c] return 0, 0 diff --git a/xrspatial/hydro/tests/test_flow_accumulation_d8.py b/xrspatial/hydro/tests/test_flow_accumulation_d8.py index e5db1c719..3f2510732 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_d8.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_d8.py @@ -626,3 +626,56 @@ def test_weight_dataset_accessor(): expected = flow_accumulation(agg, weight=w).data for var in ('a', 'b'): np.testing.assert_allclose(out[var].data, expected, equal_nan=True) + + +# --------------------------------------------------------------------------- +# D8 code -> (dy, dx) lookup (#3738) +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("code, expected", [ + # the eight valid codes: E, SE, S, SW, W, NW, N, NE + (1, (0, 1)), + (2, (1, 1)), + (4, (1, 0)), + (8, (1, -1)), + (16, (0, -1)), + (32, (-1, -1)), + (64, (-1, 0)), + (128, (-1, 1)), + # float codes as they arrive from a float64 flow-direction raster + (4.0, (1, 0)), + (128.0, (-1, 1)), + # no-flow / pit + (0, (0, 0)), + (0.0, (0, 0)), + # in-range but not a power of two + (3, (0, 0)), + (5, (0, 0)), + # outside the table + (129, (0, 0)), + (255, (0, 0)), + (-1, (0, 0)), + (1e9, (0, 0)), + (-1e9, (0, 0)), + # NaN: int(nan) is INT64_MIN inside numba, the guard must catch it + # before the table is indexed (run under NUMBA_BOUNDSCHECK=1 to check) + (np.nan, (0, 0)), +]) +def test_code_to_offset_matches_if_chain(code, expected): + from xrspatial.hydro.flow_accumulation_d8 import ( + _code_to_offset, + _code_to_offset_py, + ) + dy, dx = _code_to_offset(code) + assert (dy, dx) == expected + assert isinstance(dy, (int, np.integer)) + assert isinstance(dx, (int, np.integer)) + if code == code: # _code_to_offset_py raises on NaN like int(nan) does + assert _code_to_offset_py(code) == expected + + +def test_code_to_offset_tables_only_populate_d8_codes(): + from xrspatial.hydro.flow_accumulation_d8 import _D8_DX, _D8_DY + assert _D8_DY.shape == _D8_DX.shape == (129,) + populated = np.flatnonzero((_D8_DY != 0) | (_D8_DX != 0)) + assert populated.tolist() == [1, 2, 4, 8, 16, 32, 64, 128] From 2405f0cfec46eeea04e49ea636b743de87d49ef6 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Fri, 4 Sep 2026 11:53:48 -0400 Subject: [PATCH 2/2] Address review: explicit NaN check in _code_to_offset, module-level test imports (#3738) Return (0, 0) for a NaN code before calling int() rather than relying on the integer guard catching whatever the float-to-int conversion produced (INT64_MIN on x86, 0 on aarch64, poison to LLVM). Reword the guard comment to describe the real hazard, a rejection test on the float that is False for NaN, instead of claiming an integer rejection test would miss it. The extra compare is free: 36.5 ms and 37.1 ms on two runs of the 1000x2000 case versus 39.7 ms before it, all within noise, output still identical to origin/main. Move the private imports in the new tests up to module level to match the rest of the file. --- xrspatial/hydro/flow_accumulation_d8.py | 11 ++++++++--- xrspatial/hydro/tests/test_flow_accumulation_d8.py | 11 ++++++----- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/xrspatial/hydro/flow_accumulation_d8.py b/xrspatial/hydro/flow_accumulation_d8.py index e6e8950d9..40ede45be 100644 --- a/xrspatial/hydro/flow_accumulation_d8.py +++ b/xrspatial/hydro/flow_accumulation_d8.py @@ -210,10 +210,15 @@ def _no_weight_cupy(): @ngjit def _code_to_offset(code): """Return (dy, dx) row/col offset for a D8 direction code.""" + # NaN never reaches int(): the float-to-int conversion of NaN is + # undefined (INT64_MIN on x86, 0 on aarch64) and numba does no bounds + # checking, so the table index must never depend on it. + if code != code: + return 0, 0 c = int(code) - # Inside-the-box guard, not a rejection test: int(nan) is INT64_MIN in - # numba and there is no bounds checking, so ``c < 0 or c > 128`` style - # guards would not stop a NaN code from indexing out of bounds. + # Guard on the converted integer with an inside-the-box test. A + # rejection test on the float (``code < 0 or code > 128``) is False for + # NaN and would fall through to the lookup. if 0 <= c <= 128: return _D8_DY[c], _D8_DX[c] return 0, 0 diff --git a/xrspatial/hydro/tests/test_flow_accumulation_d8.py b/xrspatial/hydro/tests/test_flow_accumulation_d8.py index 3f2510732..66f4aff1a 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_d8.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_d8.py @@ -3,6 +3,12 @@ import xarray as xr from xrspatial.hydro import flow_accumulation +from xrspatial.hydro.flow_accumulation_d8 import ( + _D8_DX, + _D8_DY, + _code_to_offset, + _code_to_offset_py, +) from xrspatial.tests.general_checks import ( create_test_raster, cuda_and_cupy_available, @@ -662,10 +668,6 @@ def test_weight_dataset_accessor(): (np.nan, (0, 0)), ]) def test_code_to_offset_matches_if_chain(code, expected): - from xrspatial.hydro.flow_accumulation_d8 import ( - _code_to_offset, - _code_to_offset_py, - ) dy, dx = _code_to_offset(code) assert (dy, dx) == expected assert isinstance(dy, (int, np.integer)) @@ -675,7 +677,6 @@ def test_code_to_offset_matches_if_chain(code, expected): def test_code_to_offset_tables_only_populate_d8_codes(): - from xrspatial.hydro.flow_accumulation_d8 import _D8_DX, _D8_DY assert _D8_DY.shape == _D8_DX.shape == (129,) populated = np.flatnonzero((_D8_DY != 0) | (_D8_DX != 0)) assert populated.tolist() == [1, 2, 4, 8, 16, 32, 64, 128]