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

[brendancol]

Closes #3738

_code_to_offset in flow_accumulation_d8.py mapped a D8 code to its (dy, dx) 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.

• Replace the chain with two module-level int64 lookup arrays indexed by the code. numba freezes module-level arrays as compile-time constants, so the call is now a range check and two loads.
• The range check is deliberately an inside-the-box test, 0 <= c <= 128. int(nan) is INT64_MIN inside 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 erode(): NaN input causes an out-of-bounds read and floods the output with NaN #3703). Codes like 3 or 5, anything above 128, negative codes and NaN keep returning (0, 0), and the returns stay integer so r + dy in the callers is unchanged.
• _code_to_offset_py is untouched.

On a 1000x2000 float64 flow-direction raster (Gaussian bump plus default_rng(71942).normal(0, 2) noise), _flow_accum_cpu went from 56.3 ms to 39.7 ms, median of 7 after warmup, with identical results. The output is also identical on a raster built from the same DEM with 30% random NaN, checked with np.array_equal(..., equal_nan=True) against the function pulled from origin/main.

Backends: numpy and dask+numpy share this CPU kernel. cupy and dask+cupy use the device-side dispatch and are not touched.

Test plan

• New parametrised test for _code_to_offset covering the eight codes, 0, 3, 5, 129, 255, -1, 1e9, float codes and NaN, asserting the same tuples the old chain returned.
• That test run under NUMBA_BOUNDSCHECK=1, plus a negative control confirming the flag raises IndexError on an unguarded lookup with a NaN code.
• pytest xrspatial/hydro/tests -k "d8 or flow_accum or flow_length or stream_link or stream_order or watershed or flow_path or hand or basin": 872 passed.
• A/B against origin/main on the clean and 30% NaN rasters: bitwise identical.

The FlowAccumulation asv class in benchmarks/benchmarks/flow_accumulation.py already covers this path, so no benchmark change.

[brendancol]

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

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

• xrspatial/hydro/flow_accumulation_d8.py:214-216: the comment says a c < 0 or c > 128 guard would not stop a NaN code. That is not what happens here. c is already an integer at that point, and on x86 int(nan) lands on INT64_MIN, which c < 0 does reject. The hazard that actually bites is comparing the float before the conversion (code < 0 or code > 128 is False for NaN, so it falls through to int(nan)), or relying on what fptosi does with NaN at all, since LLVM treats that as poison and other targets return a different value (aarch64 fcvtzs gives 0). Reword the comment to describe that, so the next reader does not learn the wrong rule from it.
• xrspatial/hydro/flow_accumulation_d8.py:213: given the above, an explicit if code != code: return 0, 0 before int(code) would make NaN handling independent of the conversion semantics instead of depending on the integer guard catching whatever fptosi produced. The bounds-checked test proves the current x86 behaviour but not the portable one. Worth a quick timing check that the extra compare is free on the 1000x2000 case; if it is, add it.

Nits (optional improvements)

• xrspatial/hydro/tests/test_flow_accumulation_d8.py:664-668 and :677: the two new tests import _code_to_offset, _code_to_offset_py, _D8_DX and _D8_DY inside the function bodies. The rest of the file imports at module level; move these up to match.
• The PR touches the kernel behind FlowAccumulation in benchmarks/benchmarks/flow_accumulation.py, but the path-based labeler may not tag hydro as performance. Add the label by hand so the asv job runs on this PR.

What looks good

• The guard is written as 0 <= c <= 128 rather than a rejection test, and the PR includes a negative control showing NUMBA_BOUNDSCHECK=1 raises on the unguarded form with a NaN code.
• Returns stay int64 from the tables, so r + dy arithmetic in the seven importing modules is unchanged. The old literal tuples also typed as int64, so no recompilation surprises for callers.
• The A/B against origin/main covers both a clean raster and one with 30% NaN, using array_equal(equal_nan=True) rather than allclose.
• _code_to_offset_py is left alone, so the boundary-stitching code in watershed_d8, flow_length_d8 and hand_d8 that uses it is unaffected.
• The parametrised test pins the mapping explicitly (not by comparing to the Python version), so a future edit to either table is caught.

Checklist

• Algorithm matches reference (same mapping as the chain, verified for every probe)
• All implemented backends produce consistent results (numpy and dask+numpy share the kernel; GPU untouched)
• NaN handling is correct (see suggestion on the comment wording)
• Edge cases are covered by tests (0, 3, 5, 129, 255, -1, 1e9, floats, NaN)
• Dask chunk boundaries handled correctly (no change to tiling)
• No premature materialization or unnecessary copies
• Benchmark exists (FlowAccumulation)
• README feature matrix: not applicable, private helper
• Docstrings present and accurate

[brendancol]

PR Review: Replace the D8 code-to-offset if/elif chain with a lookup table (#3738), follow-up

Second pass after 2405f0c.

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

None.

Nits (optional improvements)

None.

Disposition of the first-pass findings

• Guard comment wording: fixed. flow_accumulation_d8.py:213-221 now names the real hazard (a rejection test on the float is False for NaN) instead of claiming an integer rejection test would miss it.
• Explicit NaN check before int(): fixed. if code != code: return 0, 0 sits ahead of the conversion, so the table index no longer depends on what fptosi does with NaN on a given target. Timed at 36.5 ms and 37.1 ms on the 1000x2000 case versus 39.7 ms without it, so the compare is in the noise, and the A/B against origin/main is still bitwise identical on both rasters.
• Module-level test imports: fixed. test_flow_accumulation_d8.py:6-11.
• performance label: added to the PR.

What looks good

• The NaN path is now tested three ways: the parametrised case, the same test under NUMBA_BOUNDSCHECK=1, and the A/B on a 30% NaN raster.
• The bench script's xrspatial.__file__ assertion guards against the worktree-vs-main editable install trap, so the numbers are for the branch under review.

Checklist

• Algorithm matches reference
• All implemented backends produce consistent results
• NaN handling is correct
• Edge cases are covered by tests
• Dask chunk boundaries handled correctly
• No premature materialization or unnecessary copies
• Benchmark exists
• README feature matrix: not applicable
• Docstrings present and accurate