diff --git a/xrspatial/hydro/flow_accumulation_d8.py b/xrspatial/hydro/flow_accumulation_d8.py index 677ad3ef9..40ede45be 100644 --- a/xrspatial/hydro/flow_accumulation_d8.py +++ b/xrspatial/hydro/flow_accumulation_d8.py @@ -192,26 +192,35 @@ 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.""" + # 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) - 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 + # 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 e5db1c719..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, @@ -626,3 +632,51 @@ 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): + 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(): + 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]