Skip to content

hydro: replace the D8 code-to-offset if/elif chain with a lookup table #3738

Description

@brendancol

Reason or Problem

_code_to_offset(code) in xrspatial/hydro/flow_accumulation_d8.py maps a D8 direction code to a (dy, dx) neighbour offset with an eight-way if/elif chain. _flow_accum_cpu calls it twice per cell (in-degree pass, then the BFS drain), and stream_link_d8, flow_length_d8, flow_path_d8, stream_order_d8 and hand_d8 import the same function for their own per-cell loops. On a 1000x2000 raster that is a few million data-dependent branch chains per accumulation call.

Proposal

Replace the chain with two module-level lookup arrays indexed by the code:

_D8_DY = np.zeros(129, dtype=np.int64)
_D8_DX = np.zeros(129, dtype=np.int64)
# populate entries 1, 2, 4, 8, 16, 32, 64, 128

@ngjit
def _code_to_offset(code):
    c = int(code)
    if 0 <= c <= 128:
        return _D8_DY[c], _D8_DX[c]
    return 0, 0

numba freezes module-level numpy arrays as compile-time constants, so this compiles to a range check and two loads.

Design: The guard has to be written as an inside-the-box test (0 <= c <= 128). int(nan) in numba is INT64_MIN and numba does no bounds checking, so an unguarded _D8_DY[c] on a NaN code reads out of bounds. Same bug class as the erode() fix in #3703. Non-power-of-two codes in range (3, 5, ...) and anything outside 0..128 return (0, 0) as the chain does today. Return types stay integer so r + dy in the callers is unchanged. _code_to_offset_py keeps its current behaviour.

Usage: No public API change. Modules that import _code_to_offset pick up the new version without edits.

Value: A spike on _flow_accum_cpu with a 1000x2000 D8 raster measured about 0.79x of baseline wall time with bitwise identical output.

Stakeholders and Impacts

Users of flow_accumulation, flow_length, stream_link, stream_order, flow_path and hand on the numpy and dask+numpy backends. GPU kernels have their own device-side dispatch and are not touched. The FlowAccumulation asv class in benchmarks/benchmarks/flow_accumulation.py already covers the CPU path.

Drawbacks

Two small module-level arrays in the module namespace. The guard adds a compare per call; the measurement above includes it.

Alternatives

Compacting the BFS queue seeding into the in-degree pass measured only 0.93x on its own and is left out.

Unresolved Questions

None.

Additional Notes or Context

watershed_d8.py and basin_d8.py carry private copies of _code_to_offset. Out of scope here; they can be pointed at the shared table later.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:hydrologyArea: hydrologyenhancementNew feature or requestperformancePR touches performance-sensitive code

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions