Replace the D8 code-to-offset if/elif chain with a lookup table (#3738)

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
 
 

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]