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.
Reason or Problem
_code_to_offset(code)inxrspatial/hydro/flow_accumulation_d8.pymaps a D8 direction code to a(dy, dx)neighbour offset with an eight-wayif/elifchain._flow_accum_cpucalls 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:
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 isINT64_MINand numba does no bounds checking, so an unguarded_D8_DY[c]on a NaN code reads out of bounds. Same bug class as theerode()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 sor + dyin the callers is unchanged._code_to_offset_pykeeps its current behaviour.Usage: No public API change. Modules that import
_code_to_offsetpick up the new version without edits.Value: A spike on
_flow_accum_cpuwith 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_pathandhandon the numpy and dask+numpy backends. GPU kernels have their own device-side dispatch and are not touched. TheFlowAccumulationasv class inbenchmarks/benchmarks/flow_accumulation.pyalready 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.pyandbasin_d8.pycarry private copies of_code_to_offset. Out of scope here; they can be pointed at the shared table later.