From 4ab4cc3bf47a5374df72908f09fc8c9f6d6d414e Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 26 Aug 2026 13:13:08 -0400 Subject: [PATCH 1/6] Add weight raster to flow_accumulation_d8/dinf/mfd (#3734) --- xrspatial/hydro/flow_accumulation_d8.py | 215 +++++++++++++++++----- xrspatial/hydro/flow_accumulation_dinf.py | 151 ++++++++++----- xrspatial/hydro/flow_accumulation_mfd.py | 150 +++++++++++---- 3 files changed, 396 insertions(+), 120 deletions(-) diff --git a/xrspatial/hydro/flow_accumulation_d8.py b/xrspatial/hydro/flow_accumulation_d8.py index 20bf5368c..677ad3ef9 100644 --- a/xrspatial/hydro/flow_accumulation_d8.py +++ b/xrspatial/hydro/flow_accumulation_d8.py @@ -32,8 +32,9 @@ class cupy: # type: ignore[no-redef] from xrspatial.dataset_support import supports_dataset from xrspatial.hydro._boundary_store import BoundaryStore -from xrspatial.utils import (_dask_task_name_kwargs, _validate_raster, cuda_args, - has_cuda_and_cupy, is_cupy_array, is_dask_cupy, ngjit) +from xrspatial.utils import (_dask_task_name_kwargs, _validate_matching_shape, + _validate_raster, cuda_args, has_cuda_and_cupy, + is_cupy_array, is_dask_cupy, ngjit) # ===================================================================== # Memory guards @@ -47,7 +48,9 @@ class cupy: # type: ignore[no-redef] # queue_c : int64 -> 8 # Total ~29 bytes/pixel. The caller-provided ``flow_dir`` array already # lives in RAM before the kernel runs and is not double-counted here. +# A ``weight`` raster adds one float64 copy (8 bytes/pixel). _BYTES_PER_PIXEL = 29 +_WEIGHT_BYTES_PER_PIXEL = 8 # GPU peak working set per pixel for ``_flow_accum_cupy``: # accum : float64 -> 8 @@ -89,9 +92,10 @@ def _available_gpu_memory_bytes(): return 0 -def _check_memory(height, width): +def _check_memory(height, width, weighted=False): """Raise MemoryError if the BFS kernel would exceed 50% of RAM.""" - required = int(height) * int(width) * _BYTES_PER_PIXEL + per_pixel = _BYTES_PER_PIXEL + (_WEIGHT_BYTES_PER_PIXEL if weighted else 0) + required = int(height) * int(width) * per_pixel available = _available_memory_bytes() if required > 0.5 * available: raise MemoryError( @@ -102,7 +106,7 @@ def _check_memory(height, width): ) -def _check_gpu_memory(height, width): +def _check_gpu_memory(height, width, weighted=False): """Raise MemoryError if the CuPy kernel would exceed 50% of free GPU RAM. Skips the check (returns silently) when ``_available_gpu_memory_bytes`` @@ -112,7 +116,8 @@ def _check_gpu_memory(height, width): available = _available_gpu_memory_bytes() if available <= 0: return - required = int(height) * int(width) * _GPU_BYTES_PER_PIXEL + per_pixel = _GPU_BYTES_PER_PIXEL + (_WEIGHT_BYTES_PER_PIXEL if weighted else 0) + required = int(height) * int(width) * per_pixel if required > 0.5 * available: raise MemoryError( f"flow_accumulation on a {height}x{width} grid requires " @@ -132,6 +137,57 @@ def _to_numpy_f64(arr): return np.asarray(arr, dtype=np.float64) +# ===================================================================== +# Weight helpers +# ===================================================================== +# +# Every kernel takes ``(weight, has_weight)``. When ``has_weight`` is +# 0 the ``weight`` array is a 1x1 dummy and each valid cell starts at +# 1.0 (cell count). When 1, each valid cell starts at ``weight[r, c]``, +# with NaN weights contributing 0.0 so a missing weight never alters +# the flow network or the output NaN mask (which stays ``isnan(flow_dir)``). + +_NO_WEIGHT = np.ones((1, 1), dtype=np.float64) + + +@ngjit +def _cell_weight(weight, has_weight, r, c): + """Initial accumulation for a valid cell.""" + if has_weight == 0: + return 1.0 + w = weight[r, c] + if w != w: # NaN weight -> zero contribution + return 0.0 + return w + + +def _validate_weight(weight, flow_dir, func_name): + """Validate a companion weight raster against ``flow_dir``. + + Returns the raw weight array (``weight.data``) for backend-specific + coercion by the caller. + """ + _validate_raster(weight, func_name=func_name, name='weight') + _validate_matching_shape(weight, flow_dir.shape[-2:], + func_name=func_name, name='weight', + expected_name='`flow_dir`') + return weight.data + + +def _weight_as_dask(weight_data, chunks): + """Coerce a weight array to a dask array chunked like ``flow_dir``.""" + if isinstance(weight_data, da.Array): + if weight_data.chunks != chunks: + weight_data = weight_data.rechunk(chunks) + return weight_data + return da.from_array(weight_data, chunks=chunks) + + +def _no_weight_cupy(): + import cupy as cp + return cp.ones((1, 1), dtype=cp.float64) + + # ===================================================================== # Direction helpers # ===================================================================== @@ -225,7 +281,7 @@ def _detect_flow_type(data): # ===================================================================== @ngjit -def _flow_accum_cpu(flow_dir, height, width): +def _flow_accum_cpu(flow_dir, height, width, weight, has_weight): """Kahn's BFS topological sort for flow accumulation.""" accum = np.empty((height, width), dtype=np.float64) in_degree = np.zeros((height, width), dtype=np.int32) @@ -237,7 +293,7 @@ def _flow_accum_cpu(flow_dir, height, width): v = flow_dir[r, c] if v == v: # not NaN valid[r, c] = 1 - accum[r, c] = 1.0 + accum[r, c] = _cell_weight(weight, has_weight, r, c) else: accum[r, c] = np.nan @@ -293,7 +349,8 @@ def _flow_accum_cpu(flow_dir, height, width): # ===================================================================== @cuda.jit -def _init_accum_indegree(flow_dir, accum, in_degree, state, H, W): +def _init_accum_indegree(flow_dir, accum, in_degree, state, H, W, + weight, has_weight): """Initialise accum, in_degree and state arrays on GPU.""" i, j = cuda.grid(2) if i >= H or j >= W: @@ -306,7 +363,13 @@ def _init_accum_indegree(flow_dir, accum, in_degree, state, H, W): return state[i, j] = 1 - accum[i, j] = 1.0 + if has_weight == 0: + accum[i, j] = 1.0 + else: + w = weight[i, j] + if w != w: # NaN weight -> zero contribution + w = 0.0 + accum[i, j] = w # Decode direction (inline -- can't call @ngjit from @cuda.jit) code = int(v) @@ -417,12 +480,14 @@ def _pull_from_frontier(flow_dir, accum, in_degree, state, H, W): in_degree[i, j] -= 1 -def _flow_accum_cupy(flow_dir_data): +def _flow_accum_cupy(flow_dir_data, weight=None): """GPU driver: iterative frontier peeling.""" import cupy as cp H, W = flow_dir_data.shape flow_dir_f64 = flow_dir_data.astype(cp.float64) + has_weight = 0 if weight is None else 1 + weight_f64 = _no_weight_cupy() if weight is None else weight accum = cp.zeros((H, W), dtype=cp.float64) in_degree = cp.zeros((H, W), dtype=cp.int32) @@ -432,7 +497,7 @@ def _flow_accum_cupy(flow_dir_data): griddim, blockdim = cuda_args((H, W)) _init_accum_indegree[griddim, blockdim]( - flow_dir_f64, accum, in_degree, state, H, W) + flow_dir_f64, accum, in_degree, state, H, W, weight_f64, has_weight) max_iter = H * W for _ in range(max_iter): @@ -453,7 +518,8 @@ def _flow_accum_cupy(flow_dir_data): def _flow_accum_tile_cupy(flow_dir_data, seed_top, seed_bottom, seed_left, seed_right, - seed_tl, seed_tr, seed_bl, seed_br): + seed_tl, seed_tr, seed_bl, seed_br, + weight=None): """GPU seeded flow accumulation for a single tile. Same algorithm as ``_flow_accum_cupy`` but injects external seed @@ -464,6 +530,8 @@ def _flow_accum_tile_cupy(flow_dir_data, H, W = flow_dir_data.shape flow_dir_f64 = flow_dir_data.astype(cp.float64) + has_weight = 0 if weight is None else 1 + weight_f64 = _no_weight_cupy() if weight is None else weight accum = cp.zeros((H, W), dtype=cp.float64) in_degree = cp.zeros((H, W), dtype=cp.int32) @@ -473,7 +541,7 @@ def _flow_accum_tile_cupy(flow_dir_data, griddim, blockdim = cuda_args((H, W)) _init_accum_indegree[griddim, blockdim]( - flow_dir_f64, accum, in_degree, state, H, W) + flow_dir_f64, accum, in_degree, state, H, W, weight_f64, has_weight) # Inject seeds at boundary cells. Invalid cells (state==0) are # masked to NaN at the end and never enter frontier peeling, so @@ -510,7 +578,8 @@ def _flow_accum_tile_cupy(flow_dir_data, @ngjit def _flow_accum_tile_kernel(flow_dir, h, w, seed_top, seed_bottom, seed_left, seed_right, - seed_tl, seed_tr, seed_bl, seed_br): + seed_tl, seed_tr, seed_bl, seed_br, + weight, has_weight): """Seeded BFS flow accumulation for a single tile. Same as ``_flow_accum_cpu`` but adds external seeds to boundary @@ -526,7 +595,7 @@ def _flow_accum_tile_kernel(flow_dir, h, w, v = flow_dir[r, c] if v == v: valid[r, c] = 1 - accum[r, c] = 1.0 + accum[r, c] = _cell_weight(weight, has_weight, r, c) else: accum[r, c] = np.nan @@ -734,8 +803,15 @@ def _compute_seeds(iy, ix, boundaries, flow_bdry, seed_tl, seed_tr, seed_bl, seed_br) +def _weight_tile(weight_da, iy, ix): + """Materialise one weight tile as numpy, or the dummy when unweighted.""" + if weight_da is None: + return _NO_WEIGHT, 0 + return _to_numpy_f64(weight_da.blocks[iy, ix].compute()), 1 + + def _process_tile(iy, ix, flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, weight_da=None): """Run seeded BFS on one tile; update boundaries in-place. Returns the maximum absolute boundary change (float). @@ -747,8 +823,9 @@ def _process_tile(iy, ix, flow_dir_da, boundaries, flow_bdry, seeds = _compute_seeds( iy, ix, boundaries, flow_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) + weight, has_weight = _weight_tile(weight_da, iy, ix) - accum = _flow_accum_tile_kernel(chunk, h, w, *seeds) + accum = _flow_accum_tile_kernel(chunk, h, w, *seeds, weight, has_weight) # Extract new boundary strips new_top = accum[0, :].copy() @@ -777,7 +854,7 @@ def _process_tile(iy, ix, flow_dir_da, boundaries, flow_bdry, return change -def _flow_accum_dask_iterative(flow_dir_da): +def _flow_accum_dask_iterative(flow_dir_da, weight_da=None): """Iterative boundary-propagation for arbitrarily large dask arrays. Memory usage is O(tile_size + boundary_strips) per iteration. @@ -786,6 +863,8 @@ def _flow_accum_dask_iterative(flow_dir_da): chunks_x = flow_dir_da.chunks[1] n_tile_y = len(chunks_y) n_tile_x = len(chunks_x) + if weight_da is not None: + weight_da = _weight_as_dask(weight_da, flow_dir_da.chunks) # Phase 0: extract boundary flow dirs flow_bdry = _preprocess_tiles(flow_dir_da, chunks_y, chunks_x) @@ -805,7 +884,7 @@ def _flow_accum_dask_iterative(flow_dir_da): for ix in range(n_tile_x): c = _process_tile(iy, ix, flow_dir_da, boundaries, flow_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -814,7 +893,7 @@ def _flow_accum_dask_iterative(flow_dir_da): for ix in reversed(range(n_tile_x)): c = _process_tile(iy, ix, flow_dir_da, boundaries, flow_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -826,14 +905,16 @@ def _flow_accum_dask_iterative(flow_dir_da): # Phase 3: lazy assembly via da.map_blocks return _assemble_result(flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x) + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da) def _assemble_result(flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Build a lazy dask array by re-running each tile with converged seeds.""" - def _tile_fn(flow_dir_block, block_info=None): + def _tile_fn(flow_dir_block, weight_block=None, block_info=None): if block_info is None or 0 not in block_info: return np.full(flow_dir_block.shape, np.nan, dtype=np.float64) iy, ix = block_info[0]['chunk-location'] @@ -841,20 +922,35 @@ def _tile_fn(flow_dir_block, block_info=None): seeds = _compute_seeds( iy, ix, boundaries, flow_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) + if weight_block is None: + weight, has_weight = _NO_WEIGHT, 0 + else: + weight, has_weight = _to_numpy_f64(weight_block), 1 return _flow_accum_tile_kernel( - np.asarray(flow_dir_block, dtype=np.float64), h, w, *seeds) + np.asarray(flow_dir_block, dtype=np.float64), h, w, *seeds, + weight, has_weight) + inputs = [flow_dir_da] if weight_da is None else [flow_dir_da, weight_da] return da.map_blocks( _tile_fn, - flow_dir_da, + *inputs, dtype=np.float64, meta=np.array((), dtype=np.float64), **_dask_task_name_kwargs('xrspatial.flow_accumulation_d8'), ) +def _weight_tile_cupy(weight_da, iy, ix): + """Materialise one weight tile on the GPU, or None when unweighted.""" + if weight_da is None: + return None + import cupy as cp + return cp.asarray(weight_da.blocks[iy, ix].compute(), dtype=cp.float64) + + def _process_tile_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Run seeded GPU flow accumulation on one tile; update boundaries.""" import cupy as cp @@ -865,7 +961,8 @@ def _process_tile_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry, iy, ix, boundaries, flow_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) - accum = _flow_accum_tile_cupy(chunk, *seeds) + accum = _flow_accum_tile_cupy( + chunk, *seeds, weight=_weight_tile_cupy(weight_da, iy, ix)) # Extract boundaries to CPU (small 1-D strips) new_top = accum[0, :].get().copy() @@ -893,35 +990,42 @@ def _process_tile_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry, def _assemble_result_cupy(flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Build a lazy dask+cupy array using GPU tile kernel.""" import cupy as cp - def _tile_fn(flow_dir_block, block_info=None): + def _tile_fn(flow_dir_block, weight_block=None, block_info=None): if block_info is None or 0 not in block_info: return cp.full(flow_dir_block.shape, cp.nan, dtype=cp.float64) iy, ix = block_info[0]['chunk-location'] seeds = _compute_seeds( iy, ix, boundaries, flow_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) + weight = None if weight_block is None else cp.asarray( + weight_block, dtype=cp.float64) return _flow_accum_tile_cupy( - cp.asarray(flow_dir_block, dtype=cp.float64), *seeds) + cp.asarray(flow_dir_block, dtype=cp.float64), *seeds, + weight=weight) + inputs = [flow_dir_da] if weight_da is None else [flow_dir_da, weight_da] return da.map_blocks( _tile_fn, - flow_dir_da, + *inputs, dtype=np.float64, meta=cp.array((), dtype=cp.float64), **_dask_task_name_kwargs('xrspatial.flow_accumulation_d8'), ) -def _flow_accum_dask_cupy(flow_dir_da): +def _flow_accum_dask_cupy(flow_dir_da, weight_da=None): """Dask+CuPy D8: native GPU processing per tile.""" chunks_y = flow_dir_da.chunks[0] chunks_x = flow_dir_da.chunks[1] n_tile_y = len(chunks_y) n_tile_x = len(chunks_x) + if weight_da is not None: + weight_da = _weight_as_dask(weight_da, flow_dir_da.chunks) flow_bdry = _preprocess_tiles(flow_dir_da, chunks_y, chunks_x) flow_bdry = flow_bdry.snapshot() @@ -937,7 +1041,7 @@ def _flow_accum_dask_cupy(flow_dir_da): for ix in range(n_tile_x): c = _process_tile_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -945,7 +1049,7 @@ def _flow_accum_dask_cupy(flow_dir_da): for ix in reversed(range(n_tile_x)): c = _process_tile_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -955,7 +1059,8 @@ def _flow_accum_dask_cupy(flow_dir_da): boundaries = boundaries.snapshot() return _assemble_result_cupy(flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x) + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da) # ===================================================================== @@ -964,6 +1069,7 @@ def _flow_accum_dask_cupy(flow_dir_da): @supports_dataset def flow_accumulation_d8(flow_dir: xr.DataArray, + weight: xr.DataArray | None = None, name: str = 'flow_accumulation') -> xr.DataArray: """Compute flow accumulation from a D8 flow direction grid. @@ -971,6 +1077,10 @@ def flow_accumulation_d8(flow_dir: xr.DataArray, integer D8 direction codes. For D-infinity (continuous angle) grids, use ``flow_accumulation_dinf`` instead. + By default each cell contributes 1, so the result is a count of + upstream cells. Pass ``weight`` (for example a precipitation or + melt field) to accumulate that quantity instead. + Parameters ---------- flow_dir : xarray.DataArray or xr.Dataset @@ -980,6 +1090,13 @@ def flow_accumulation_d8(flow_dir: xr.DataArray, CuPy-backed Dask. If a Dataset is passed, the operation is applied to each data variable independently. + weight : xarray.DataArray, optional + 2D raster on the same grid as ``flow_dir`` giving each cell's + own contribution. When given, each output cell is the sum of + ``weight`` over itself and every upstream cell draining + through it. NaN weights contribute 0 and do not change the + output NaN mask; weights at cells where ``flow_dir`` is NaN + are ignored. name : str, default='flow_accumulation' Name of output DataArray. @@ -988,7 +1105,9 @@ def flow_accumulation_d8(flow_dir: xr.DataArray, xarray.DataArray or xr.Dataset 2D float64 array of flow accumulation values. Each cell contains the count of upstream cells (including itself) that - drain through it. Cells with NaN flow direction produce NaN. + drain through it, or the sum of ``weight`` over those cells + when ``weight`` is given. Cells with NaN flow direction + produce NaN. References ---------- @@ -1000,17 +1119,27 @@ def flow_accumulation_d8(flow_dir: xr.DataArray, _validate_raster(flow_dir, func_name='flow_accumulation', name='flow_dir') data = flow_dir.data + weighted = weight is not None + w_data = _validate_weight(weight, flow_dir, 'flow_accumulation') \ + if weighted else None if isinstance(data, np.ndarray): - _check_memory(*data.shape) - out = _flow_accum_cpu(data.astype(np.float64), *data.shape) + _check_memory(*data.shape, weighted=weighted) + if weighted: + w_np, has_w = _to_numpy_f64(w_data), 1 + else: + w_np, has_w = _NO_WEIGHT, 0 + out = _flow_accum_cpu(data.astype(np.float64), *data.shape, + w_np, has_w) elif has_cuda_and_cupy() and is_cupy_array(data): - _check_gpu_memory(*data.shape) - out = _flow_accum_cupy(data) + import cupy as cp + _check_gpu_memory(*data.shape, weighted=weighted) + w_cp = cp.asarray(w_data, dtype=cp.float64) if weighted else None + out = _flow_accum_cupy(data, w_cp) elif has_cuda_and_cupy() and is_dask_cupy(flow_dir): - out = _flow_accum_dask_cupy(data) + out = _flow_accum_dask_cupy(data, w_data) elif da is not None and isinstance(data, da.Array): - out = _flow_accum_dask_iterative(data) + out = _flow_accum_dask_iterative(data, w_data) else: raise TypeError(f"Unsupported array type: {type(data)}") diff --git a/xrspatial/hydro/flow_accumulation_dinf.py b/xrspatial/hydro/flow_accumulation_dinf.py index 769f3cb0d..c629ce850 100644 --- a/xrspatial/hydro/flow_accumulation_dinf.py +++ b/xrspatial/hydro/flow_accumulation_dinf.py @@ -41,8 +41,16 @@ class cupy: # type: ignore[no-redef] from xrspatial.hydro._boundary_store import BoundaryStore from xrspatial.dataset_support import supports_dataset from xrspatial.hydro.flow_accumulation_d8 import ( + _NO_WEIGHT, + _WEIGHT_BYTES_PER_PIXEL, + _cell_weight, _find_ready_and_finalize, + _no_weight_cupy, _preprocess_tiles, + _validate_weight, + _weight_as_dask, + _weight_tile, + _weight_tile_cupy, ) @@ -103,9 +111,10 @@ def _available_gpu_memory_bytes(): return 0 -def _check_memory(height, width): +def _check_memory(height, width, weighted=False): """Raise MemoryError if the BFS kernel would exceed 50% of RAM.""" - required = int(height) * int(width) * _BYTES_PER_PIXEL + per_pixel = _BYTES_PER_PIXEL + (_WEIGHT_BYTES_PER_PIXEL if weighted else 0) + required = int(height) * int(width) * per_pixel available = _available_memory_bytes() if required > 0.5 * available: raise MemoryError( @@ -116,7 +125,7 @@ def _check_memory(height, width): ) -def _check_gpu_memory(height, width): +def _check_gpu_memory(height, width, weighted=False): """Raise MemoryError if the CuPy kernel would exceed 50% of free GPU RAM. Skips the check (returns silently) when ``_available_gpu_memory_bytes`` @@ -126,7 +135,8 @@ def _check_gpu_memory(height, width): available = _available_gpu_memory_bytes() if available <= 0: return - required = int(height) * int(width) * _GPU_BYTES_PER_PIXEL + per_pixel = _GPU_BYTES_PER_PIXEL + (_WEIGHT_BYTES_PER_PIXEL if weighted else 0) + required = int(height) * int(width) * per_pixel if required > 0.5 * available: raise MemoryError( f"flow_accumulation_dinf on a {height}x{width} grid requires " @@ -201,7 +211,7 @@ def _angle_to_neighbors(angle): # ===================================================================== @ngjit -def _flow_accum_dinf_cpu(flow_dir, height, width): +def _flow_accum_dinf_cpu(flow_dir, height, width, weight, has_weight): """Kahn's BFS topological sort for Dinf flow accumulation.""" accum = np.empty((height, width), dtype=np.float64) in_degree = np.zeros((height, width), dtype=np.int32) @@ -212,7 +222,7 @@ def _flow_accum_dinf_cpu(flow_dir, height, width): v = flow_dir[r, c] if v == v: # not NaN valid[r, c] = 1 - accum[r, c] = 1.0 + accum[r, c] = _cell_weight(weight, has_weight, r, c) else: accum[r, c] = np.nan @@ -277,7 +287,8 @@ def _flow_accum_dinf_cpu(flow_dir, height, width): # ===================================================================== @cuda.jit -def _init_accum_indegree_dinf(flow_dir, accum, in_degree, state, H, W): +def _init_accum_indegree_dinf(flow_dir, accum, in_degree, state, H, W, + weight, has_weight): """Initialise accum/in_degree/state for Dinf on GPU.""" i, j = cuda.grid(2) if i >= H or j >= W: @@ -290,7 +301,13 @@ def _init_accum_indegree_dinf(flow_dir, accum, in_degree, state, H, W): return state[i, j] = 1 - accum[i, j] = 1.0 + if has_weight == 0: + accum[i, j] = 1.0 + else: + w = weight[i, j] + if w != w: # NaN weight -> zero contribution + w = 0.0 + accum[i, j] = w if v < 0.0: # pit return @@ -421,12 +438,14 @@ def _pull_from_frontier_dinf(flow_dir, accum, in_degree, state, H, W): in_degree[i, j] -= 1 -def _flow_accum_dinf_cupy(flow_dir_data): +def _flow_accum_dinf_cupy(flow_dir_data, weight=None): """GPU driver: iterative frontier peeling for Dinf.""" import cupy as cp H, W = flow_dir_data.shape flow_dir_f64 = flow_dir_data.astype(cp.float64) + has_weight = 0 if weight is None else 1 + weight_f64 = _no_weight_cupy() if weight is None else weight accum = cp.zeros((H, W), dtype=cp.float64) in_degree = cp.zeros((H, W), dtype=cp.int32) @@ -436,7 +455,7 @@ def _flow_accum_dinf_cupy(flow_dir_data): griddim, blockdim = cuda_args((H, W)) _init_accum_indegree_dinf[griddim, blockdim]( - flow_dir_f64, accum, in_degree, state, H, W) + flow_dir_f64, accum, in_degree, state, H, W, weight_f64, has_weight) max_iter = H * W for _ in range(max_iter): @@ -456,12 +475,15 @@ def _flow_accum_dinf_cupy(flow_dir_data): def _flow_accum_dinf_tile_cupy(flow_dir_data, seed_top, seed_bottom, seed_left, seed_right, - seed_tl, seed_tr, seed_bl, seed_br): + seed_tl, seed_tr, seed_bl, seed_br, + weight=None): """GPU seeded Dinf flow accumulation for a single tile.""" import cupy as cp H, W = flow_dir_data.shape flow_dir_f64 = flow_dir_data.astype(cp.float64) + has_weight = 0 if weight is None else 1 + weight_f64 = _no_weight_cupy() if weight is None else weight accum = cp.zeros((H, W), dtype=cp.float64) in_degree = cp.zeros((H, W), dtype=cp.int32) @@ -471,7 +493,7 @@ def _flow_accum_dinf_tile_cupy(flow_dir_data, griddim, blockdim = cuda_args((H, W)) _init_accum_indegree_dinf[griddim, blockdim]( - flow_dir_f64, accum, in_degree, state, H, W) + flow_dir_f64, accum, in_degree, state, H, W, weight_f64, has_weight) accum[0, :] += cp.asarray(seed_top) accum[H - 1, :] += cp.asarray(seed_bottom) @@ -506,7 +528,8 @@ def _flow_accum_dinf_tile_cupy(flow_dir_data, def _flow_accum_dinf_tile_kernel(flow_dir, h, w, seed_top, seed_bottom, seed_left, seed_right, - seed_tl, seed_tr, seed_bl, seed_br): + seed_tl, seed_tr, seed_bl, seed_br, + weight, has_weight): """Seeded BFS Dinf flow accumulation for a single tile.""" accum = np.empty((h, w), dtype=np.float64) in_degree = np.zeros((h, w), dtype=np.int32) @@ -517,7 +540,7 @@ def _flow_accum_dinf_tile_kernel(flow_dir, h, w, v = flow_dir[r, c] if v == v: valid[r, c] = 1 - accum[r, c] = 1.0 + accum[r, c] = _cell_weight(weight, has_weight, r, c) else: accum[r, c] = np.nan @@ -723,7 +746,8 @@ def _compute_seeds_dinf(iy, ix, boundaries, flow_bdry, def _process_tile_dinf(iy, ix, flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Run seeded Dinf BFS on one tile; update boundaries in-place.""" chunk = np.asarray( flow_dir_da.blocks[iy, ix].compute(), dtype=np.float64) @@ -732,8 +756,10 @@ def _process_tile_dinf(iy, ix, flow_dir_da, boundaries, flow_bdry, seeds = _compute_seeds_dinf( iy, ix, boundaries, flow_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) + weight, has_weight = _weight_tile(weight_da, iy, ix) - accum = _flow_accum_dinf_tile_kernel(chunk, h, w, *seeds) + accum = _flow_accum_dinf_tile_kernel(chunk, h, w, *seeds, + weight, has_weight) new_top = accum[0, :].copy() new_bottom = accum[-1, :].copy() @@ -759,12 +785,14 @@ def _process_tile_dinf(iy, ix, flow_dir_da, boundaries, flow_bdry, return change -def _flow_accum_dinf_dask_iterative(flow_dir_da): +def _flow_accum_dinf_dask_iterative(flow_dir_da, weight_da=None): """Iterative boundary-propagation for Dinf dask arrays.""" chunks_y = flow_dir_da.chunks[0] chunks_x = flow_dir_da.chunks[1] n_tile_y = len(chunks_y) n_tile_x = len(chunks_x) + if weight_da is not None: + weight_da = _weight_as_dask(weight_da, flow_dir_da.chunks) flow_bdry = _preprocess_tiles(flow_dir_da, chunks_y, chunks_x) flow_bdry = flow_bdry.snapshot() @@ -780,7 +808,7 @@ def _flow_accum_dinf_dask_iterative(flow_dir_da): for ix in range(n_tile_x): c = _process_tile_dinf(iy, ix, flow_dir_da, boundaries, flow_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -788,7 +816,7 @@ def _flow_accum_dinf_dask_iterative(flow_dir_da): for ix in reversed(range(n_tile_x)): c = _process_tile_dinf(iy, ix, flow_dir_da, boundaries, flow_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -798,14 +826,16 @@ def _flow_accum_dinf_dask_iterative(flow_dir_da): boundaries = boundaries.snapshot() return _assemble_result_dinf(flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x) + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da) def _assemble_result_dinf(flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Build lazy dask array by re-running each Dinf tile with converged seeds.""" - def _tile_fn(flow_dir_block, block_info=None): + def _tile_fn(flow_dir_block, weight_block=None, block_info=None): if block_info is None or 0 not in block_info: return np.full(flow_dir_block.shape, np.nan, dtype=np.float64) iy, ix = block_info[0]['chunk-location'] @@ -813,12 +843,18 @@ def _tile_fn(flow_dir_block, block_info=None): seeds = _compute_seeds_dinf( iy, ix, boundaries, flow_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) + if weight_block is None: + weight, has_weight = _NO_WEIGHT, 0 + else: + weight, has_weight = _to_numpy_f64(weight_block), 1 return _flow_accum_dinf_tile_kernel( - np.asarray(flow_dir_block, dtype=np.float64), h, w, *seeds) + np.asarray(flow_dir_block, dtype=np.float64), h, w, *seeds, + weight, has_weight) + inputs = [flow_dir_da] if weight_da is None else [flow_dir_da, weight_da] return da.map_blocks( _tile_fn, - flow_dir_da, + *inputs, dtype=np.float64, meta=np.array((), dtype=np.float64), **_dask_task_name_kwargs('xrspatial.flow_accumulation_dinf'), @@ -826,7 +862,8 @@ def _tile_fn(flow_dir_block, block_info=None): def _process_tile_dinf_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Run seeded GPU Dinf flow accumulation on one tile.""" import cupy as cp @@ -837,7 +874,8 @@ def _process_tile_dinf_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry, iy, ix, boundaries, flow_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) - accum = _flow_accum_dinf_tile_cupy(chunk, *seeds) + accum = _flow_accum_dinf_tile_cupy( + chunk, *seeds, weight=_weight_tile_cupy(weight_da, iy, ix)) new_top = accum[0, :].get().copy() new_bottom = accum[-1, :].get().copy() @@ -864,35 +902,42 @@ def _process_tile_dinf_cupy(iy, ix, flow_dir_da, boundaries, flow_bdry, def _assemble_result_dinf_cupy(flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Build lazy dask+cupy array using GPU Dinf tile kernel.""" import cupy as cp - def _tile_fn(flow_dir_block, block_info=None): + def _tile_fn(flow_dir_block, weight_block=None, block_info=None): if block_info is None or 0 not in block_info: return cp.full(flow_dir_block.shape, cp.nan, dtype=cp.float64) iy, ix = block_info[0]['chunk-location'] seeds = _compute_seeds_dinf( iy, ix, boundaries, flow_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) + weight = None if weight_block is None else cp.asarray( + weight_block, dtype=cp.float64) return _flow_accum_dinf_tile_cupy( - cp.asarray(flow_dir_block, dtype=cp.float64), *seeds) + cp.asarray(flow_dir_block, dtype=cp.float64), *seeds, + weight=weight) + inputs = [flow_dir_da] if weight_da is None else [flow_dir_da, weight_da] return da.map_blocks( _tile_fn, - flow_dir_da, + *inputs, dtype=np.float64, meta=cp.array((), dtype=cp.float64), **_dask_task_name_kwargs('xrspatial.flow_accumulation_dinf'), ) -def _flow_accum_dinf_dask_cupy(flow_dir_da): +def _flow_accum_dinf_dask_cupy(flow_dir_da, weight_da=None): """Dask+CuPy Dinf: native GPU processing per tile.""" chunks_y = flow_dir_da.chunks[0] chunks_x = flow_dir_da.chunks[1] n_tile_y = len(chunks_y) n_tile_x = len(chunks_x) + if weight_da is not None: + weight_da = _weight_as_dask(weight_da, flow_dir_da.chunks) flow_bdry = _preprocess_tiles(flow_dir_da, chunks_y, chunks_x) flow_bdry = flow_bdry.snapshot() @@ -909,7 +954,7 @@ def _flow_accum_dinf_dask_cupy(flow_dir_da): c = _process_tile_dinf_cupy( iy, ix, flow_dir_da, boundaries, flow_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -918,7 +963,7 @@ def _flow_accum_dinf_dask_cupy(flow_dir_da): c = _process_tile_dinf_cupy( iy, ix, flow_dir_da, boundaries, flow_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -928,7 +973,8 @@ def _flow_accum_dinf_dask_cupy(flow_dir_da): boundaries = boundaries.snapshot() return _assemble_result_dinf_cupy(flow_dir_da, boundaries, flow_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x) + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da) # ===================================================================== @@ -937,6 +983,7 @@ def _flow_accum_dinf_dask_cupy(flow_dir_da): @supports_dataset def flow_accumulation_dinf(flow_dir: xr.DataArray, + weight: xr.DataArray | None = None, name: str = 'flow_accumulation') -> xr.DataArray: """Compute flow accumulation from a D-infinity flow direction grid. @@ -945,6 +992,10 @@ def flow_accumulation_dinf(flow_dir: xr.DataArray, area. Flow is split proportionally between two neighbors following Tarboton (1997). + By default each cell contributes 1, so the result is a (fractional) + count of upstream cells. Pass ``weight`` (for example a + precipitation or melt field) to accumulate that quantity instead. + Parameters ---------- flow_dir : xarray.DataArray or xr.Dataset @@ -955,6 +1006,13 @@ def flow_accumulation_dinf(flow_dir: xr.DataArray, CuPy-backed Dask. If a Dataset is passed, the operation is applied to each data variable independently. + weight : xarray.DataArray, optional + 2-D raster on the same grid as ``flow_dir`` giving each cell's + own contribution. When given, each output cell is the + proportionally split sum of ``weight`` over itself and every + upstream cell draining through it. NaN weights contribute 0 + and do not change the output NaN mask; weights at cells where + ``flow_dir`` is NaN are ignored. name : str, default='flow_accumulation' Name of output DataArray. @@ -964,7 +1022,8 @@ def flow_accumulation_dinf(flow_dir: xr.DataArray, 2-D float64 array of flow accumulation values. Each cell holds the total upstream contributing area (including itself) that drains through it, weighted by D-inf proportional - splitting. NaN where the input has NaN. + splitting, or the split sum of ``weight`` over those cells + when ``weight`` is given. NaN where the input has NaN. References ---------- @@ -976,17 +1035,27 @@ def flow_accumulation_dinf(flow_dir: xr.DataArray, name='flow_dir') data = flow_dir.data + weighted = weight is not None + w_data = _validate_weight(weight, flow_dir, 'flow_accumulation_dinf') \ + if weighted else None if isinstance(data, np.ndarray): - _check_memory(*data.shape) - out = _flow_accum_dinf_cpu(data.astype(np.float64), *data.shape) + _check_memory(*data.shape, weighted=weighted) + if weighted: + w_np, has_w = _to_numpy_f64(w_data), 1 + else: + w_np, has_w = _NO_WEIGHT, 0 + out = _flow_accum_dinf_cpu(data.astype(np.float64), *data.shape, + w_np, has_w) elif has_cuda_and_cupy() and is_cupy_array(data): - _check_gpu_memory(*data.shape) - out = _flow_accum_dinf_cupy(data) + import cupy as cp + _check_gpu_memory(*data.shape, weighted=weighted) + w_cp = cp.asarray(w_data, dtype=cp.float64) if weighted else None + out = _flow_accum_dinf_cupy(data, w_cp) elif has_cuda_and_cupy() and is_dask_cupy(flow_dir): - out = _flow_accum_dinf_dask_cupy(data) + out = _flow_accum_dinf_dask_cupy(data, w_data) elif da is not None and isinstance(data, da.Array): - out = _flow_accum_dinf_dask_iterative(data) + out = _flow_accum_dinf_dask_iterative(data, w_data) else: raise TypeError(f"Unsupported array type: {type(data)}") diff --git a/xrspatial/hydro/flow_accumulation_mfd.py b/xrspatial/hydro/flow_accumulation_mfd.py index d3db24d2e..e954deca6 100644 --- a/xrspatial/hydro/flow_accumulation_mfd.py +++ b/xrspatial/hydro/flow_accumulation_mfd.py @@ -41,6 +41,14 @@ class cupy: # type: ignore[no-redef] ) from xrspatial.hydro._boundary_store import BoundaryStore from xrspatial.dataset_support import supports_dataset +from xrspatial.hydro.flow_accumulation_d8 import ( + _NO_WEIGHT, + _WEIGHT_BYTES_PER_PIXEL, + _cell_weight, + _no_weight_cupy, + _validate_weight, + _weight_as_dask, +) # ===================================================================== @@ -98,9 +106,10 @@ def _available_gpu_memory_bytes(): return 0 -def _check_memory(height, width): +def _check_memory(height, width, weighted=False): """Raise MemoryError if the BFS kernel would exceed 50% of RAM.""" - required = int(height) * int(width) * _BYTES_PER_PIXEL + per_pixel = _BYTES_PER_PIXEL + (_WEIGHT_BYTES_PER_PIXEL if weighted else 0) + required = int(height) * int(width) * per_pixel available = _available_memory_bytes() if required > 0.5 * available: raise MemoryError( @@ -111,7 +120,7 @@ def _check_memory(height, width): ) -def _check_gpu_memory(height, width): +def _check_gpu_memory(height, width, weighted=False): """Raise MemoryError if the CuPy kernel would exceed 50% of free GPU RAM. Skips the check (returns silently) when ``_available_gpu_memory_bytes`` @@ -121,7 +130,8 @@ def _check_gpu_memory(height, width): available = _available_gpu_memory_bytes() if available <= 0: return - required = int(height) * int(width) * _GPU_BYTES_PER_PIXEL + per_pixel = _GPU_BYTES_PER_PIXEL + (_WEIGHT_BYTES_PER_PIXEL if weighted else 0) + required = int(height) * int(width) * per_pixel if required > 0.5 * available: raise MemoryError( f"flow_accumulation_mfd on a {height}x{width} grid requires " @@ -145,13 +155,15 @@ def _check_gpu_memory(height, width): # ===================================================================== @ngjit -def _flow_accum_mfd_cpu(fractions, height, width): +def _flow_accum_mfd_cpu(fractions, height, width, weight, has_weight): """Kahn's BFS topological sort for MFD flow accumulation. Parameters ---------- fractions : (8, H, W) float64 array of flow fractions height, width : int + weight : (H, W) float64 per-cell contribution (1x1 dummy when unused) + has_weight : int, 1 to use ``weight`` instead of a unit count Returns ------- @@ -173,7 +185,7 @@ def _flow_accum_mfd_cpu(fractions, height, width): accum[r, c] = np.nan else: valid[r, c] = 1 - accum[r, c] = 1.0 + accum[r, c] = _cell_weight(weight, has_weight, r, c) n_valid += 1 # Pass 2: compute in-degrees @@ -238,7 +250,8 @@ def _flow_accum_mfd_cpu(fractions, height, width): # ===================================================================== @cuda.jit -def _init_accum_indegree_mfd(fractions, accum, in_degree, state, H, W): +def _init_accum_indegree_mfd(fractions, accum, in_degree, state, H, W, + weight, has_weight): """Initialise accum, in_degree and state for MFD on GPU.""" i, j = cuda.grid(2) if i >= H or j >= W: @@ -251,7 +264,13 @@ def _init_accum_indegree_mfd(fractions, accum, in_degree, state, H, W): return state[i, j] = 1 - accum[i, j] = 1.0 + if has_weight == 0: + accum[i, j] = 1.0 + else: + w = weight[i, j] + if w != w: # NaN weight -> zero contribution + w = 0.0 + accum[i, j] = w # Neighbor offsets: E, SE, S, SW, W, NW, N, NE for k in range(8): @@ -359,12 +378,14 @@ def _pull_from_frontier_mfd(fractions, accum, in_degree, state, H, W): in_degree[i, j] -= 1 -def _flow_accum_mfd_cupy(fractions_data): +def _flow_accum_mfd_cupy(fractions_data, weight=None): """GPU driver: iterative frontier peeling for MFD.""" import cupy as cp _, H, W = fractions_data.shape fractions_f64 = fractions_data.astype(cp.float64) + has_weight = 0 if weight is None else 1 + weight_f64 = _no_weight_cupy() if weight is None else weight accum = cp.zeros((H, W), dtype=cp.float64) in_degree = cp.zeros((H, W), dtype=cp.int32) @@ -374,7 +395,7 @@ def _flow_accum_mfd_cupy(fractions_data): griddim, blockdim = cuda_args((H, W)) _init_accum_indegree_mfd[griddim, blockdim]( - fractions_f64, accum, in_degree, state, H, W) + fractions_f64, accum, in_degree, state, H, W, weight_f64, has_weight) max_iter = H * W for _ in range(max_iter): @@ -400,12 +421,15 @@ def _flow_accum_mfd_cupy(fractions_data): def _flow_accum_mfd_tile_kernel(fractions, h, w, seed_top, seed_bottom, seed_left, seed_right, - seed_tl, seed_tr, seed_bl, seed_br): + seed_tl, seed_tr, seed_bl, seed_br, + weight, has_weight): """Seeded BFS MFD flow accumulation for a single tile. Parameters ---------- fractions : (8, h, w) float64 -- MFD flow fractions for this tile + weight : (h, w) float64 per-cell contribution (1x1 dummy when unused) + has_weight : int, 1 to use ``weight`` instead of a unit count """ dy = np.array([0, 1, 1, 1, 0, -1, -1, -1], dtype=np.int64) dx = np.array([1, 1, 0, -1, -1, -1, 0, 1], dtype=np.int64) @@ -421,7 +445,7 @@ def _flow_accum_mfd_tile_kernel(fractions, h, w, v = fractions[0, r, c] if v == v: # not NaN valid[r, c] = 1 - accum[r, c] = 1.0 + accum[r, c] = _cell_weight(weight, has_weight, r, c) n_valid += 1 else: accum[r, c] = np.nan @@ -666,7 +690,8 @@ def _compute_seeds_mfd(iy, ix, boundaries, frac_bdry, def _process_tile_mfd(iy, ix, fractions_da, boundaries, frac_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Run seeded MFD BFS on one tile; update boundaries in-place.""" # Extract this tile's fractions: (8, tile_h, tile_w) y_start = sum(chunks_y[:iy]) @@ -683,7 +708,16 @@ def _process_tile_mfd(iy, ix, fractions_da, boundaries, frac_bdry, iy, ix, boundaries, frac_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) - accum = _flow_accum_mfd_tile_kernel(chunk, h, w, *seeds) + if weight_da is None: + weight, has_weight = _NO_WEIGHT, 0 + else: + weight = np.asarray( + weight_da[y_start:y_end, x_start:x_end].compute(), + dtype=np.float64) + has_weight = 1 + + accum = _flow_accum_mfd_tile_kernel(chunk, h, w, *seeds, + weight, has_weight) # NaN cells don't contribute flow; replace with 0 for boundary storage new_top = np.where(np.isnan(accum[0, :]), 0.0, accum[0, :]) @@ -710,16 +744,20 @@ def _process_tile_mfd(iy, ix, fractions_da, boundaries, frac_bdry, return change -def _flow_accum_mfd_dask_iterative(fractions_da, chunks_y, chunks_x): +def _flow_accum_mfd_dask_iterative(fractions_da, chunks_y, chunks_x, + weight_da=None): """Iterative boundary-propagation for MFD dask arrays. Parameters ---------- fractions_da : dask array of shape (8, H, W) chunks_y, chunks_x : tuples of chunk sizes for the spatial dims + weight_da : optional (H, W) array of per-cell contributions """ n_tile_y = len(chunks_y) n_tile_x = len(chunks_x) + if weight_da is not None: + weight_da = _weight_as_dask(weight_da, (chunks_y, chunks_x)) # The 8 direction bands must stay in a single chunk: every tile kernel # needs all 8 fractions, and the lazy assembly drops axis 0 per block. @@ -742,7 +780,7 @@ def _flow_accum_mfd_dask_iterative(fractions_da, chunks_y, chunks_x): for ix in range(n_tile_x): c = _process_tile_mfd(iy, ix, fractions_da, boundaries, frac_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -750,7 +788,7 @@ def _flow_accum_mfd_dask_iterative(fractions_da, chunks_y, chunks_x): for ix in reversed(range(n_tile_x)): c = _process_tile_mfd(iy, ix, fractions_da, boundaries, frac_bdry, chunks_y, chunks_x, - n_tile_y, n_tile_x) + n_tile_y, n_tile_x, weight_da) if c > max_change: max_change = c @@ -760,11 +798,13 @@ def _flow_accum_mfd_dask_iterative(fractions_da, chunks_y, chunks_x): boundaries = boundaries.snapshot() return _assemble_result_mfd(fractions_da, boundaries, frac_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x) + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da) def _assemble_result_mfd(fractions_da, boundaries, frac_bdry, - chunks_y, chunks_x, n_tile_y, n_tile_x): + chunks_y, chunks_x, n_tile_y, n_tile_x, + weight_da=None): """Build a lazy dask array by re-running each MFD tile with converged seeds. fractions_da is (8, H, W) chunked one tile per (chunks_y, chunks_x) @@ -777,7 +817,7 @@ def _assemble_result_mfd(fractions_da, boundaries, frac_bdry, y_starts = np.cumsum((0,) + tuple(chunks_y[:-1])) x_starts = np.cumsum((0,) + tuple(chunks_x[:-1])) - def _tile(chunk, block_info=None): + def _tile(chunk, weight_block=None, block_info=None): # block_info[0]['array-location'] gives ((0, 8), (y0, y1), (x0, x1)). loc = block_info[0]['array-location'] y0 = loc[1][0] @@ -790,16 +830,26 @@ def _tile(chunk, block_info=None): seeds = _compute_seeds_mfd( iy, ix, boundaries, frac_bdry, chunks_y, chunks_x, n_tile_y, n_tile_x) - return _flow_accum_mfd_tile_kernel(chunk, h, w, *seeds) - + if weight_block is None: + weight, has_weight = _NO_WEIGHT, 0 + else: + weight = np.asarray(weight_block, dtype=np.float64) + has_weight = 1 + return _flow_accum_mfd_tile_kernel(chunk, h, w, *seeds, + weight, has_weight) + + # The 2-D weight aligns with the trailing (y, x) axes of the 3-D + # fractions array, so map_blocks pairs each spatial tile correctly. + inputs = [fractions_da] if weight_da is None else [fractions_da, weight_da] return da.map_blocks( - _tile, fractions_da, drop_axis=0, + _tile, *inputs, drop_axis=0, dtype=np.float64, meta=np.array((), dtype=np.float64), **_dask_task_name_kwargs('xrspatial.flow_accumulation_mfd'), ) -def _flow_accum_mfd_dask_cupy(fractions_da, chunks_y, chunks_x): +def _flow_accum_mfd_dask_cupy(fractions_da, chunks_y, chunks_x, + weight_da=None): """Dask+CuPy MFD: convert to numpy, run iterative, convert back.""" import cupy as cp @@ -807,7 +857,14 @@ def _flow_accum_mfd_dask_cupy(fractions_da, chunks_y, chunks_x): lambda b: b.get(), dtype=fractions_da.dtype, meta=np.array((), dtype=fractions_da.dtype), ) - result = _flow_accum_mfd_dask_iterative(fractions_np, chunks_y, chunks_x) + if weight_da is not None: + weight_da = _weight_as_dask(weight_da, (chunks_y, chunks_x)) + weight_da = weight_da.map_blocks( + lambda b: b.get() if hasattr(b, 'get') else b, + dtype=weight_da.dtype, meta=np.array((), dtype=weight_da.dtype), + ) + result = _flow_accum_mfd_dask_iterative(fractions_np, chunks_y, chunks_x, + weight_da) return result.map_blocks( cp.asarray, dtype=result.dtype, meta=cp.array((), dtype=result.dtype), @@ -820,14 +877,15 @@ def _flow_accum_mfd_dask_cupy(fractions_da, chunks_y, chunks_x): @supports_dataset def flow_accumulation_mfd(flow_dir_mfd: xr.DataArray, + weight: xr.DataArray | None = None, name: str = 'flow_accumulation_mfd') -> xr.DataArray: """Compute flow accumulation from an MFD flow direction grid. Takes the 3-D fractional output of ``flow_direction_mfd`` and accumulates upstream contributing area through all downslope - paths simultaneously. Each cell starts with a value of 1 (itself) - and passes fractions of its accumulated value to each downstream - neighbor. + paths simultaneously. Each cell starts with a value of 1 (itself), + or its ``weight`` when given, and passes fractions of its + accumulated value to each downstream neighbor. Parameters ---------- @@ -840,6 +898,13 @@ def flow_accumulation_mfd(flow_dir_mfd: xr.DataArray, CuPy-backed Dask. If a Dataset is passed, the operation is applied to each data variable independently. + weight : xarray.DataArray, optional + 2-D ``(H, W)`` raster on the same grid as ``flow_dir_mfd`` + giving each cell's own contribution (for example precipitation + or melt). When given, each output cell is the fraction-split + sum of ``weight`` over itself and every upstream cell draining + through it. NaN weights contribute 0 and do not change the + output NaN mask; weights at nodata cells are ignored. name : str, default='flow_accumulation_mfd' Name of output DataArray. @@ -848,8 +913,9 @@ def flow_accumulation_mfd(flow_dir_mfd: xr.DataArray, xarray.DataArray or xr.Dataset 2-D float64 array of flow accumulation values. Each cell holds the total upstream contributing area (including itself) - that drains through it, weighted by MFD fractions. - NaN where the input has NaN. + that drains through it, weighted by MFD fractions, or the + split sum of ``weight`` over those cells when ``weight`` is + given. NaN where the input has NaN. References ---------- @@ -877,22 +943,34 @@ def flow_accumulation_mfd(flow_dir_mfd: xr.DataArray, _validate_mfd_fractions(data, func_name='flow_accumulation_mfd', name='flow_dir_mfd') + weighted = weight is not None + w_data = _validate_weight(weight, flow_dir_mfd, 'flow_accumulation_mfd') \ + if weighted else None + if isinstance(data, np.ndarray): - _check_memory(data.shape[1], data.shape[2]) + _check_memory(data.shape[1], data.shape[2], weighted=weighted) + if weighted: + w_np = w_data.get() if hasattr(w_data, 'get') else w_data + w_np, has_w = np.asarray(w_np, dtype=np.float64), 1 + else: + w_np, has_w = _NO_WEIGHT, 0 out = _flow_accum_mfd_cpu( - data.astype(np.float64), data.shape[1], data.shape[2]) + data.astype(np.float64), data.shape[1], data.shape[2], + w_np, has_w) elif has_cuda_and_cupy() and is_cupy_array(data): - _check_gpu_memory(data.shape[1], data.shape[2]) - out = _flow_accum_mfd_cupy(data) + import cupy as cp + _check_gpu_memory(data.shape[1], data.shape[2], weighted=weighted) + w_cp = cp.asarray(w_data, dtype=cp.float64) if weighted else None + out = _flow_accum_mfd_cupy(data, w_cp) elif has_cuda_and_cupy() and is_dask_cupy(flow_dir_mfd): # Spatial chunk sizes from dims 1 and 2 chunks_y = data.chunks[1] chunks_x = data.chunks[2] - out = _flow_accum_mfd_dask_cupy(data, chunks_y, chunks_x) + out = _flow_accum_mfd_dask_cupy(data, chunks_y, chunks_x, w_data) elif da is not None and isinstance(data, da.Array): chunks_y = data.chunks[1] chunks_x = data.chunks[2] - out = _flow_accum_mfd_dask_iterative(data, chunks_y, chunks_x) + out = _flow_accum_mfd_dask_iterative(data, chunks_y, chunks_x, w_data) else: raise TypeError(f"Unsupported array type: {type(data)}") From 146baa12117ac9fd3bb9dcdafc5d2d1a84494cff Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 26 Aug 2026 13:17:14 -0400 Subject: [PATCH 2/6] Test weighted flow accumulation across backends (#3734) --- .../hydro/tests/test_flow_accumulation_d8.py | 133 ++++++++++++++++++ .../tests/test_flow_accumulation_dinf.py | 118 ++++++++++++++++ .../hydro/tests/test_flow_accumulation_mfd.py | 81 +++++++++++ 3 files changed, 332 insertions(+) diff --git a/xrspatial/hydro/tests/test_flow_accumulation_d8.py b/xrspatial/hydro/tests/test_flow_accumulation_d8.py index 778e6e168..3acc31e04 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_d8.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_d8.py @@ -446,3 +446,136 @@ def test_degenerate_shape(shape): assert result.shape == shape assert not np.isnan(result.data).any() np.testing.assert_array_equal(result.data, 1.0) + + +# --------------------------------------------------------------------------- +# Weighted accumulation (#3734) +# --------------------------------------------------------------------------- + +_BOWL_FLOW_DIR = np.array([ + [np.nan, np.nan, np.nan, np.nan, np.nan], + [np.nan, 2, 2, 4, np.nan], + [np.nan, 2, 2, 4, np.nan], + [np.nan, 1, 1, 0, np.nan], + [np.nan, np.nan, np.nan, np.nan, np.nan], +], dtype=np.float64) + + +def test_weight_known_values(): + """Each cell is the sum of weight over itself and its upstream cells.""" + flow_dir = np.array([ + [1.0, 1.0, 1.0, 0.0], + [1.0, 1.0, 1.0, 0.0], + ], dtype=np.float64) + weight = np.array([ + [1.0, 2.0, 3.0, 4.0], + [0.5, 0.0, -1.0, 2.0], + ], dtype=np.float64) + expected = np.array([ + [1.0, 3.0, 6.0, 10.0], + [0.5, 0.5, -0.5, 1.5], + ]) + agg = create_test_raster(flow_dir) + w = create_test_raster(weight) + result = flow_accumulation(agg, weight=w) + np.testing.assert_allclose(result.data, expected) + + +def test_weight_ones_equals_count(): + """weight=1 everywhere reproduces the unweighted cell count.""" + agg = create_test_raster(_BOWL_FLOW_DIR) + w = create_test_raster(np.ones_like(_BOWL_FLOW_DIR)) + np.testing.assert_allclose( + flow_accumulation(agg, weight=w).data, + flow_accumulation(agg).data, equal_nan=True) + + +def test_weight_nan_contributes_zero(): + """NaN weight at a valid cell adds 0; the NaN mask still follows flow_dir.""" + flow_dir = np.array([ + [1.0, 1.0, 1.0, 0.0], + [np.nan, 1.0, 1.0, 0.0], + ], dtype=np.float64) + weight = np.array([ + [1.0, np.nan, 3.0, 4.0], + [7.0, 2.0, np.nan, 1.0], + ], dtype=np.float64) + expected = np.array([ + [1.0, 1.0, 4.0, 8.0], + [np.nan, 2.0, 2.0, 3.0], + ]) + agg = create_test_raster(flow_dir) + w = create_test_raster(weight) + result = flow_accumulation(agg, weight=w) + np.testing.assert_allclose(result.data, expected, equal_nan=True) + assert np.array_equal(np.isnan(result.data), np.isnan(flow_dir)) + + +def test_weight_shape_mismatch_raises(): + agg = create_test_raster(_BOWL_FLOW_DIR) + w = create_test_raster(np.ones((5, 4))) + with pytest.raises(ValueError, match="weight"): + flow_accumulation(agg, weight=w) + + +def test_weight_not_dataarray_raises(): + agg = create_test_raster(_BOWL_FLOW_DIR) + with pytest.raises(TypeError, match="weight"): + flow_accumulation(agg, weight=np.ones((5, 5))) + + +def _weighted_reference(): + flow_dir = _make_cross_backend_flow_dir() + rng = np.random.default_rng(3734) + weight = rng.random(flow_dir.shape) * 10.0 + weight[2, 3] = np.nan + weight[5, 1] = -4.0 + np_result = flow_accumulation( + create_test_raster(flow_dir), weight=create_test_raster(weight)) + return flow_dir, weight, np_result.data + + +@dask_array_available +@pytest.mark.parametrize("chunks", [(3, 3), (5, 5), (2, 6), (1, 1), (6, 6)]) +def test_weight_dask_equals_numpy(chunks): + """Weighted dask matches numpy for every chunk layout.""" + flow_dir, weight, expected = _weighted_reference() + dask_agg = create_test_raster(flow_dir, backend='dask', chunks=chunks) + dask_w = create_test_raster(weight, backend='dask', chunks=chunks) + np.testing.assert_allclose( + flow_accumulation(dask_agg, weight=dask_w).data.compute(), + expected, equal_nan=True) + + +@dask_array_available +def test_weight_dask_rechunks_weight(): + """A numpy weight or one with different chunks is aligned to flow_dir.""" + flow_dir, weight, expected = _weighted_reference() + dask_agg = create_test_raster(flow_dir, backend='dask', chunks=(3, 3)) + np.testing.assert_allclose( + flow_accumulation(dask_agg, weight=create_test_raster(weight)) + .data.compute(), expected, equal_nan=True) + dask_w = create_test_raster(weight, backend='dask', chunks=(2, 5)) + np.testing.assert_allclose( + flow_accumulation(dask_agg, weight=dask_w).data.compute(), + expected, equal_nan=True) + + +@cuda_and_cupy_available +def test_weight_cupy_equals_numpy(): + flow_dir, weight, expected = _weighted_reference() + cupy_agg = create_test_raster(flow_dir, backend='cupy') + cupy_w = create_test_raster(weight, backend='cupy') + np.testing.assert_allclose( + flow_accumulation(cupy_agg, weight=cupy_w).data.get(), + expected, equal_nan=True) + + +@cuda_and_cupy_available +def test_weight_dask_cupy_equals_numpy(): + flow_dir, weight, expected = _weighted_reference() + agg = create_test_raster(flow_dir, backend='dask+cupy', chunks=(3, 3)) + w = create_test_raster(weight, backend='dask+cupy', chunks=(3, 3)) + np.testing.assert_allclose( + flow_accumulation(agg, weight=w).data.compute().get(), + expected, equal_nan=True) diff --git a/xrspatial/hydro/tests/test_flow_accumulation_dinf.py b/xrspatial/hydro/tests/test_flow_accumulation_dinf.py index 3b84a2f0e..2f8bfe46a 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_dinf.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_dinf.py @@ -272,3 +272,121 @@ def test_error_message_mentions_dimensions(self): ): with pytest.raises(MemoryError, match=r"7x11"): flow_accumulation_dinf(agg) + + +# =========================================================================== +# Weighted accumulation (#3734) +# =========================================================================== + +def _dinf_weighted_case(): + """3x3 with the centre pit fed by cardinal + diagonal neighbours.""" + pi = np.pi + flow_dir = np.array([ + [7 * pi / 4, 3 * pi / 2, 5 * pi / 4], + [0.0, -1.0, pi], + [pi / 4, pi / 2, 3 * pi / 4], + ], dtype=np.float64) + weight = np.array([ + [1.0, 2.0, 3.0], + [4.0, 10.0, np.nan], + [0.5, 1.5, -2.0], + ], dtype=np.float64) + return flow_dir, weight + + +def test_dinf_weight_known_values(): + """Centre pit collects the weights of all 8 neighbours plus its own.""" + flow_dir, weight = _dinf_weighted_case() + result = flow_accumulation_dinf( + create_test_raster(flow_dir), weight=create_test_raster(weight)) + # Each neighbour keeps its own weight (nothing upstream of it); + # the NaN weight contributes 0 at (1, 2). + expected = weight.copy() + expected[1, 2] = 0.0 + expected[1, 1] = 10.0 + 1 + 2 + 3 + 4 + 0 + 0.5 + 1.5 - 2 + np.testing.assert_allclose(result.data, expected) + assert not np.isnan(result.data).any() + + +def test_dinf_weight_ones_equals_count(): + flow_dir, _ = _dinf_weighted_case() + agg = create_test_raster(flow_dir) + np.testing.assert_allclose( + flow_accumulation_dinf(agg, weight=create_test_raster( + np.ones_like(flow_dir))).data, + flow_accumulation_dinf(agg).data, equal_nan=True) + + +def test_dinf_weight_split_proportionally(): + """A weight flowing at pi/4 is split 50/50 between E and NE... via D-inf.""" + pi = np.pi + # Bottom-left cell flows at pi/8: half-way between E (0) and NE (pi/4). + flow_dir = np.array([ + [-1.0, -1.0], + [pi / 8, -1.0], + ], dtype=np.float64) + weight = np.array([[0.0, 0.0], [8.0, 0.0]], dtype=np.float64) + result = flow_accumulation_dinf( + create_test_raster(flow_dir), weight=create_test_raster(weight)) + np.testing.assert_allclose(result.data[1, 0], 8.0) + np.testing.assert_allclose(result.data[1, 1], 4.0) # E + np.testing.assert_allclose(result.data[0, 1], 4.0) # NE + np.testing.assert_allclose(result.data[0, 0], 0.0) + + +def test_dinf_weight_shape_mismatch_raises(): + flow_dir, weight = _dinf_weighted_case() + with pytest.raises(ValueError, match="weight"): + flow_accumulation_dinf(create_test_raster(flow_dir), + weight=create_test_raster(weight[:, :2])) + + +def _dinf_weighted_reference(): + """Acyclic D-inf grid derived from a noisy tilted DEM. + + Random angles would form cycles, which the BFS does not define + consistently across tilings, so derive directions from a surface. + """ + from xrspatial.hydro import flow_direction_dinf + rng = np.random.default_rng(3734) + elev = rng.random((12, 15)) + np.add.outer( + np.arange(12) * 0.5, np.arange(15) * 0.3) + elev[6, 7] = np.nan + flow_dir = flow_direction_dinf(create_test_raster(elev)).data + weight = rng.random(flow_dir.shape) * 5.0 + weight[3, 4] = np.nan + expected = flow_accumulation_dinf( + create_test_raster(flow_dir), weight=create_test_raster(weight)).data + return flow_dir, weight, expected + + +@dask_array_available +@pytest.mark.parametrize("chunks", [(3, 3), (4, 5), (12, 15), (1, 1)]) +def test_dinf_weight_dask_equals_numpy(chunks): + flow_dir, weight, expected = _dinf_weighted_reference() + dask_agg = create_test_raster(flow_dir, backend='dask', chunks=chunks) + dask_w = create_test_raster(weight, backend='dask', chunks=(2, 7)) + np.testing.assert_allclose( + flow_accumulation_dinf(dask_agg, weight=dask_w).data.compute(), + expected, equal_nan=True) + + +@cuda_and_cupy_available +def test_dinf_weight_cupy_equals_numpy(): + flow_dir, weight, expected = _dinf_weighted_reference() + np.testing.assert_allclose( + flow_accumulation_dinf( + create_test_raster(flow_dir, backend='cupy'), + weight=create_test_raster(weight, backend='cupy')).data.get(), + expected, equal_nan=True) + + +@cuda_and_cupy_available +def test_dinf_weight_dask_cupy_equals_numpy(): + flow_dir, weight, expected = _dinf_weighted_reference() + np.testing.assert_allclose( + flow_accumulation_dinf( + create_test_raster(flow_dir, backend='dask+cupy', chunks=(4, 5)), + weight=create_test_raster(weight, backend='dask+cupy', + chunks=(4, 5))).data.compute().get(), + expected, equal_nan=True) diff --git a/xrspatial/hydro/tests/test_flow_accumulation_mfd.py b/xrspatial/hydro/tests/test_flow_accumulation_mfd.py index d8af7b1cc..6bdbb0bb1 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_mfd.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_mfd.py @@ -507,3 +507,84 @@ def test_dask_cycle_raises(self): mfd = _make_cyclic_mfd(backend='dask') with pytest.raises(ValueError, match="cycle"): flow_accumulation_mfd(mfd).data.compute() + + +class TestFlowAccumulationMFDWeight: + """Weighted accumulation (#3734).""" + + @staticmethod + def _weighted_case(): + elev = _make_plane_south(7, 7) + mfd = flow_direction_mfd(elev) + rng = np.random.default_rng(3734) + weight = xr.DataArray(rng.random((7, 7)) * 3.0, dims=['y', 'x']) + weight.values[2, 3] = np.nan + return mfd, weight + + def test_ones_equals_count(self): + mfd, weight = self._weighted_case() + np.testing.assert_allclose( + flow_accumulation_mfd(mfd, weight=xr.ones_like(weight)).values, + flow_accumulation_mfd(mfd).values, equal_nan=True) + + def test_mass_conserved_on_plane(self): + """On a south-sloping plane every interior column drains to the + bottom interior row, so the bottom row sums to the sum of all + finite interior weights that do not leave through the side edges.""" + elev = _make_plane_south(7, 7) + mfd = flow_direction_mfd(elev) + # Uniform weight on a plane flowing straight south: bottom-row value + # at column c is the column sum of interior weights above it. + weight = xr.DataArray(np.full((7, 7), 2.5), dims=['y', 'x']) + accum = flow_accumulation_mfd(mfd, weight=weight) + count = flow_accumulation_mfd(mfd) + np.testing.assert_allclose(accum.values, count.values * 2.5, + equal_nan=True) + + def test_nan_weight_contributes_zero(self): + mfd, weight = self._weighted_case() + filled = weight.fillna(0.0) + np.testing.assert_allclose( + flow_accumulation_mfd(mfd, weight=weight).values, + flow_accumulation_mfd(mfd, weight=filled).values, equal_nan=True) + # Mask still follows the fraction grid, not the weight. + accum = flow_accumulation_mfd(mfd, weight=weight) + assert np.array_equal(np.isnan(accum.values), + np.isnan(mfd.values[0])) + + def test_shape_mismatch_raises(self): + mfd, weight = self._weighted_case() + with pytest.raises(ValueError, match="weight"): + flow_accumulation_mfd(mfd, weight=weight[:, :5]) + + @pytest.mark.parametrize("chunks", [(3, 3), (7, 7), (2, 5)]) + def test_dask_equals_numpy(self, chunks): + pytest.importorskip("dask.array") + mfd, weight = self._weighted_case() + expected = flow_accumulation_mfd(mfd, weight=weight).values + mfd_dask = mfd.chunk({'neighbor': 8, 'y': chunks[0], 'x': chunks[1]}) + for w in (weight, weight.chunk({'y': 4, 'x': 2})): + np.testing.assert_allclose( + flow_accumulation_mfd(mfd_dask, weight=w).values, + expected, equal_nan=True) + + def test_cupy_equals_numpy(self): + from xrspatial.utils import has_cuda_and_cupy + if not has_cuda_and_cupy(): + pytest.skip("CUDA/cupy not available") + import cupy + import dask.array as da + mfd, weight = self._weighted_case() + expected = flow_accumulation_mfd(mfd, weight=weight).values + mfd_cp = mfd.copy(data=cupy.asarray(mfd.values)) + w_cp = weight.copy(data=cupy.asarray(weight.values)) + np.testing.assert_allclose( + flow_accumulation_mfd(mfd_cp, weight=w_cp).data.get(), + expected, equal_nan=True) + mfd_dc = mfd.copy(data=da.from_array(cupy.asarray(mfd.values), + chunks=(8, 3, 3))) + w_dc = weight.copy(data=da.from_array(cupy.asarray(weight.values), + chunks=(3, 3))) + np.testing.assert_allclose( + flow_accumulation_mfd(mfd_dc, weight=w_dc).data.compute().get(), + expected, equal_nan=True) From 0916b099b0195ae24f144f06774776095f91ee82 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 26 Aug 2026 13:18:17 -0400 Subject: [PATCH 3/6] Show weighted flow accumulation in the hydrology user guide (#3734) --- examples/user_guide/11_Hydrology.ipynb | 44 ++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/examples/user_guide/11_Hydrology.ipynb b/examples/user_guide/11_Hydrology.ipynb index 1812e21a4..0e8b16a59 100644 --- a/examples/user_guide/11_Hydrology.ipynb +++ b/examples/user_guide/11_Hydrology.ipynb @@ -416,6 +416,50 @@ "ax.set_axis_off()" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Weighted accumulation\n", + "\n", + "Counting cells assumes every cell contributes the same amount of water. Passing a `weight` raster on the same grid changes the quantity being accumulated: each output cell becomes the sum of `weight` over itself and everything upstream of it. With a precipitation or snowmelt field as the weight, the result is an accumulated flux rather than a contributing-cell count.\n", + "\n", + "The example below builds a synthetic orographic precipitation field (wetter at higher elevation and toward the west) and compares the weighted result with the plain count. NaN cells in `weight` contribute zero and do not change the output NaN mask, which still follows `flow_dir`. Fill NaN explicitly if you want a different convention.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Synthetic precipitation (mm): rises with elevation and toward the west.\n", + "elev_norm = (dem - float(np.nanmin(dem.values))) / float(np.ptp(dem.values[~np.isnan(dem.values)]))\n", + "west_east = xr.DataArray(np.linspace(1.0, 0.4, W)[None, :].repeat(H, axis=0),\n", + " dims=dem.dims, coords=dem.coords)\n", + "precip = (400 + 1200 * elev_norm) * west_east\n", + "precip.name = 'precip'\n", + "\n", + "precip_accum = xrspatial.flow_accumulation(flow_dir, weight=precip)\n", + "\n", + "print(f\"Cell count at outlet: {np.nanmax(flow_accum.values):.0f}\")\n", + "print(f\"Accumulated precipitation at outlet: {np.nanmax(precip_accum.values):.3g} mm-cells\")\n", + "\n", + "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", + "fig.patch.set_facecolor('white')\n", + "for ax, data, label in zip(axes, [flow_accum, precip_accum],\n", + " ['Upstream cell count', 'Accumulated precipitation']):\n", + " ax.set_facecolor('white')\n", + " data.plot.imshow(ax=ax, cmap=water_cmap,\n", + " norm=LogNorm(vmin=float(np.nanmin(data.values[data.values > 0])),\n", + " vmax=float(np.nanmax(data.values))),\n", + " add_colorbar=True, cbar_kwargs={'label': f'{label} (log scale)',\n", + " 'shrink': 0.7})\n", + " ax.set_title(label)\n", + " ax.set_axis_off()\n", + "plt.tight_layout()\n" + ] + }, { "cell_type": "markdown", "id": "httgn03vcco", From 1d802e0cb6da927467ebcd8248abdbf27b992559 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 26 Aug 2026 13:24:45 -0400 Subject: [PATCH 4/6] Address review: numpy-coerce cupy weights in MFD dask path, add dtype/Dataset/forwarding tests (#3734) --- xrspatial/hydro/flow_accumulation_mfd.py | 11 +++--- .../hydro/tests/test_flow_accumulation_d8.py | 35 +++++++++++++++++++ .../tests/test_flow_accumulation_dinf.py | 11 +++--- .../hydro/tests/test_flow_accumulation_mfd.py | 9 ++--- 4 files changed, 49 insertions(+), 17 deletions(-) diff --git a/xrspatial/hydro/flow_accumulation_mfd.py b/xrspatial/hydro/flow_accumulation_mfd.py index e954deca6..8c51bbf79 100644 --- a/xrspatial/hydro/flow_accumulation_mfd.py +++ b/xrspatial/hydro/flow_accumulation_mfd.py @@ -46,6 +46,7 @@ class cupy: # type: ignore[no-redef] _WEIGHT_BYTES_PER_PIXEL, _cell_weight, _no_weight_cupy, + _to_numpy_f64, _validate_weight, _weight_as_dask, ) @@ -711,9 +712,8 @@ def _process_tile_mfd(iy, ix, fractions_da, boundaries, frac_bdry, if weight_da is None: weight, has_weight = _NO_WEIGHT, 0 else: - weight = np.asarray( - weight_da[y_start:y_end, x_start:x_end].compute(), - dtype=np.float64) + weight = _to_numpy_f64( + weight_da[y_start:y_end, x_start:x_end].compute()) has_weight = 1 accum = _flow_accum_mfd_tile_kernel(chunk, h, w, *seeds, @@ -833,7 +833,7 @@ def _tile(chunk, weight_block=None, block_info=None): if weight_block is None: weight, has_weight = _NO_WEIGHT, 0 else: - weight = np.asarray(weight_block, dtype=np.float64) + weight = _to_numpy_f64(weight_block) has_weight = 1 return _flow_accum_mfd_tile_kernel(chunk, h, w, *seeds, weight, has_weight) @@ -950,8 +950,7 @@ def flow_accumulation_mfd(flow_dir_mfd: xr.DataArray, if isinstance(data, np.ndarray): _check_memory(data.shape[1], data.shape[2], weighted=weighted) if weighted: - w_np = w_data.get() if hasattr(w_data, 'get') else w_data - w_np, has_w = np.asarray(w_np, dtype=np.float64), 1 + w_np, has_w = _to_numpy_f64(w_data), 1 else: w_np, has_w = _NO_WEIGHT, 0 out = _flow_accum_mfd_cpu( diff --git a/xrspatial/hydro/tests/test_flow_accumulation_d8.py b/xrspatial/hydro/tests/test_flow_accumulation_d8.py index 3acc31e04..a2ad19259 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_d8.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_d8.py @@ -579,3 +579,38 @@ def test_weight_dask_cupy_equals_numpy(): np.testing.assert_allclose( flow_accumulation(agg, weight=w).data.compute().get(), expected, equal_nan=True) + + +def test_weight_integer_dtype(): + """Integer weights are coerced to float64.""" + agg = create_test_raster(_BOWL_FLOW_DIR) + w_int = create_test_raster(np.full((5, 5), 3, dtype=np.int32)) + w_float = create_test_raster(np.full((5, 5), 3.0)) + result = flow_accumulation(agg, weight=w_int) + assert result.dtype == np.float64 + np.testing.assert_allclose( + result.data, flow_accumulation(agg, weight=w_float).data, equal_nan=True) + + +def test_weight_dataset_input(): + """A Dataset flow_dir applies the same weight to every variable.""" + agg = create_test_raster(_BOWL_FLOW_DIR) + w = create_test_raster(np.full((5, 5), 2.0)) + ds = xr.Dataset({'a': agg, 'b': agg}) + out = flow_accumulation(ds, weight=w) + expected = flow_accumulation(agg, weight=w).data + for var in ('a', 'b'): + np.testing.assert_allclose(out[var].data, expected, equal_nan=True) + + +def test_weight_forwarded_by_accessor_and_routing(): + """weight= reaches the router via flow_accumulation(routing=) and .xrs.""" + import xrspatial.accessor # noqa: F401 + agg = create_test_raster(_BOWL_FLOW_DIR) + w = create_test_raster(np.full((5, 5), 2.0)) + expected = flow_accumulation(agg, weight=w).data + np.testing.assert_allclose( + flow_accumulation(agg, weight=w, routing='d8').data, expected, + equal_nan=True) + np.testing.assert_allclose( + agg.xrs.flow_accumulation(weight=w).data, expected, equal_nan=True) diff --git a/xrspatial/hydro/tests/test_flow_accumulation_dinf.py b/xrspatial/hydro/tests/test_flow_accumulation_dinf.py index 2f8bfe46a..97145e28f 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_dinf.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_dinf.py @@ -318,7 +318,7 @@ def test_dinf_weight_ones_equals_count(): def test_dinf_weight_split_proportionally(): - """A weight flowing at pi/4 is split 50/50 between E and NE... via D-inf.""" + """A weight flowing at pi/8 is split 50/50 between E and NE by D-inf.""" pi = np.pi # Bottom-left cell flows at pi/8: half-way between E (0) and NE (pi/4). flow_dir = np.array([ @@ -365,10 +365,11 @@ def _dinf_weighted_reference(): def test_dinf_weight_dask_equals_numpy(chunks): flow_dir, weight, expected = _dinf_weighted_reference() dask_agg = create_test_raster(flow_dir, backend='dask', chunks=chunks) - dask_w = create_test_raster(weight, backend='dask', chunks=(2, 7)) - np.testing.assert_allclose( - flow_accumulation_dinf(dask_agg, weight=dask_w).data.compute(), - expected, equal_nan=True) + for w_chunks in (chunks, (2, 7)): + dask_w = create_test_raster(weight, backend='dask', chunks=w_chunks) + np.testing.assert_allclose( + flow_accumulation_dinf(dask_agg, weight=dask_w).data.compute(), + expected, equal_nan=True) @cuda_and_cupy_available diff --git a/xrspatial/hydro/tests/test_flow_accumulation_mfd.py b/xrspatial/hydro/tests/test_flow_accumulation_mfd.py index 6bdbb0bb1..4947558dc 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_mfd.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_mfd.py @@ -527,14 +527,11 @@ def test_ones_equals_count(self): flow_accumulation_mfd(mfd, weight=xr.ones_like(weight)).values, flow_accumulation_mfd(mfd).values, equal_nan=True) - def test_mass_conserved_on_plane(self): - """On a south-sloping plane every interior column drains to the - bottom interior row, so the bottom row sums to the sum of all - finite interior weights that do not leave through the side edges.""" + def test_uniform_weight_scales_count(self): + """A uniform weight w scales the unweighted count by w everywhere, + since the MFD split is linear in the accumulated value.""" elev = _make_plane_south(7, 7) mfd = flow_direction_mfd(elev) - # Uniform weight on a plane flowing straight south: bottom-row value - # at column c is the column sum of interior weights above it. weight = xr.DataArray(np.full((7, 7), 2.5), dims=['y', 'x']) accum = flow_accumulation_mfd(mfd, weight=weight) count = flow_accumulation_mfd(mfd) From 868b2b100f1e04e7998854ae62dccab7a76d9468 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 26 Aug 2026 13:47:06 -0400 Subject: [PATCH 5/6] Test weight through the Dataset accessor (#3734) --- xrspatial/hydro/tests/test_flow_accumulation_d8.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/xrspatial/hydro/tests/test_flow_accumulation_d8.py b/xrspatial/hydro/tests/test_flow_accumulation_d8.py index a2ad19259..e5db1c719 100644 --- a/xrspatial/hydro/tests/test_flow_accumulation_d8.py +++ b/xrspatial/hydro/tests/test_flow_accumulation_d8.py @@ -614,3 +614,15 @@ def test_weight_forwarded_by_accessor_and_routing(): equal_nan=True) np.testing.assert_allclose( agg.xrs.flow_accumulation(weight=w).data, expected, equal_nan=True) + + +def test_weight_dataset_accessor(): + """ds.xrs.flow_accumulation(weight=) weights every variable.""" + import xrspatial.accessor # noqa: F401 + agg = create_test_raster(_BOWL_FLOW_DIR) + w = create_test_raster(np.full((5, 5), 2.0)) + ds = xr.Dataset({'a': agg, 'b': agg}) + out = ds.xrs.flow_accumulation(weight=w) + expected = flow_accumulation(agg, weight=w).data + for var in ('a', 'b'): + np.testing.assert_allclose(out[var].data, expected, equal_nan=True) From 990e5eddc74aad2e189fbe05f58bc5cf246fa81a Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Wed, 26 Aug 2026 21:23:18 -0400 Subject: [PATCH 6/6] Add physical accuracy checks for weighted flow accumulation (#3734) --- .../test_flow_accumulation_weight_accuracy.py | 311 ++++++++++++++++++ 1 file changed, 311 insertions(+) create mode 100644 xrspatial/hydro/tests/test_flow_accumulation_weight_accuracy.py diff --git a/xrspatial/hydro/tests/test_flow_accumulation_weight_accuracy.py b/xrspatial/hydro/tests/test_flow_accumulation_weight_accuracy.py new file mode 100644 index 000000000..330b6efad --- /dev/null +++ b/xrspatial/hydro/tests/test_flow_accumulation_weight_accuracy.py @@ -0,0 +1,311 @@ +"""Accuracy checks for weighted flow accumulation (#3734). + +These tests do not compare against hand-built grids. They check physical +identities that any correct weighted accumulation must satisfy on a real +terrain surface: + +1. Mass balance: every unit of weight entering the network leaves through + exactly one outlet (a pit) or leaks off the valid grid where a cell + drains into a NaN neighbour. The leak term is computed from + ``flow_dir`` alone, independently of the accumulation kernels. +2. Watershed cross-check (D8): the accumulation at a pour point equals + the weight integrated over the watershed traced *downstream* to that + point, which is a separate code path from the upstream BFS. +3. Linearity: accumulation is a linear operator in ``weight``. + +The weight is an orographic "melt" field (zero below a snowline, growing +with elevation above it) rather than random noise, so the sign and scale +of any error show up in the identities. +""" + +import numpy as np +import pytest +import xarray as xr + +import xrspatial.accessor # noqa: F401 +from xrspatial.hydro import ( + basin, + flow_accumulation, + flow_direction_d8, + flow_direction_dinf, + flow_direction_mfd, + watershed, +) +from xrspatial.hydro.flow_accumulation_d8 import _code_to_offset_py +from xrspatial.hydro.flow_accumulation_dinf import _angle_to_neighbors +from xrspatial.tests.general_checks import ( + cuda_and_cupy_available, + dask_array_available, +) + +_N = 80 +_CHUNKS = (25, 25) # ragged on an 80x80 grid + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture(scope='module') +def dem(): + grid = xr.DataArray(np.zeros((_N, _N)), dims=['y', 'x']) + return grid.xrs.generate_terrain(seed=3734) + + +@pytest.fixture(scope='module') +def melt(dem): + """Orographic melt: 0 below the snowline, linear in elevation above.""" + e = dem.values + lo, hi = np.nanmin(e), np.nanmax(e) + snowline = lo + 0.4 * (hi - lo) + rate = np.clip((e - snowline) / (hi - snowline), 0.0, None) * 12.0 + return xr.DataArray(rate, dims=dem.dims, coords=dem.coords, name='melt') + + +@pytest.fixture(scope='module') +def glacier(dem): + """Indicator of glacierised surface: the top 30% of the elevation range.""" + e = dem.values + lo, hi = np.nanmin(e), np.nanmax(e) + return xr.DataArray((e > lo + 0.7 * (hi - lo)).astype(np.float64), + dims=dem.dims, coords=dem.coords, name='glacier') + + +@pytest.fixture(scope='module') +def routers(dem): + return { + 'd8': flow_direction_d8(dem), + 'dinf': flow_direction_dinf(dem), + 'mfd': flow_direction_mfd(dem), + } + + +def _to_backend(agg, backend): + """Move a numpy-backed DataArray to *backend*.""" + if backend == 'numpy': + return agg + if backend == 'dask': + return agg.chunk({'y': _CHUNKS[0], 'x': _CHUNKS[1]}) + import cupy + out = agg.copy(data=cupy.asarray(agg.values)) + if backend == 'dask+cupy': + import dask.array as da + # Keep the 8 MFD bands in one chunk; tile only the spatial axes. + chunks = ((agg.shape[0],) if agg.ndim == 3 else ()) + _CHUNKS + out = agg.copy(data=da.from_array(cupy.asarray(agg.values), + chunks=chunks)) + return out + + +def _to_numpy(agg): + data = agg.data + if hasattr(data, 'compute'): + data = data.compute() + if hasattr(data, 'get'): + data = data.get() + return np.asarray(data) + + +_BACKENDS = [ + 'numpy', + pytest.param('dask', marks=dask_array_available), + pytest.param('cupy', marks=cuda_and_cupy_available), + pytest.param('dask+cupy', marks=cuda_and_cupy_available), +] + + +# --------------------------------------------------------------------------- +# Leak computation (independent of the accumulation kernels) +# --------------------------------------------------------------------------- + +def _valid_at(valid, r, c): + h, w = valid.shape + return 0 <= r < h and 0 <= c < w and valid[r, c] + + +def _outflow_loss_fraction(routing, fdir): + """Per-cell fraction of outflow that drains into NaN or off the grid. + + Pits have no outflow and are handled separately as outlets. + """ + if routing == 'mfd': + frac = fdir # (8, H, W) + valid = ~np.isnan(frac[0]) + dy = [0, 1, 1, 1, 0, -1, -1, -1] + dx = [1, 1, 0, -1, -1, -1, 0, 1] + else: + valid = ~np.isnan(fdir) + h, w = valid.shape + loss = np.zeros((h, w), dtype=np.float64) + for r in range(h): + for c in range(w): + if not valid[r, c]: + continue + if routing == 'd8': + oy, ox = _code_to_offset_py(fdir[r, c]) + if (oy, ox) != (0, 0) and not _valid_at(valid, r + oy, c + ox): + loss[r, c] = 1.0 + elif routing == 'dinf': + dy1, dx1, w1, dy2, dx2, w2 = _angle_to_neighbors(fdir[r, c]) + if w1 > 0 and not _valid_at(valid, r + dy1, c + dx1): + loss[r, c] += w1 + if w2 > 0 and not _valid_at(valid, r + dy2, c + dx2): + loss[r, c] += w2 + else: + for k in range(8): + f = frac[k, r, c] + if f > 0 and not _valid_at(valid, r + dy[k], c + dx[k]): + loss[r, c] += f + return loss + + +def _pit_mask(routing, fdir): + if routing == 'd8': + return fdir == 0 + if routing == 'dinf': + return fdir == -1.0 + valid = ~np.isnan(fdir[0]) + return valid & (np.nansum(fdir, axis=0) == 0) + + +# --------------------------------------------------------------------------- +# 1. Mass balance +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize('routing', ['d8', 'dinf', 'mfd']) +@pytest.mark.parametrize('backend', _BACKENDS) +def test_mass_balance(routers, melt, routing, backend): + """sum(weight) == sum(accum at pits) + sum(accum * leak fraction).""" + fdir = routers[routing] + fd_np = fdir.values + accum = _to_numpy(flow_accumulation( + _to_backend(fdir, backend), weight=_to_backend(melt, backend), + routing=routing)) + + valid = ~np.isnan(fd_np[0] if routing == 'mfd' else fd_np) + assert np.array_equal(np.isnan(accum), ~valid) + + total_in = float(np.sum(melt.values[valid])) + pits = _pit_mask(routing, fd_np) + loss = _outflow_loss_fraction(routing, fd_np) + total_out = float(np.sum(accum[pits])) + float( + np.nansum(accum * loss)) + assert total_in > 0 + np.testing.assert_allclose(total_out, total_in, rtol=1e-9) + + +def test_mass_balance_per_basin(routers, melt): + """D8: weight summed over each basin equals the accumulation at its outlet.""" + fdir = routers['d8'] + accum = flow_accumulation(fdir, weight=melt).values + basins = basin(fdir).values + ids = np.unique(basins[~np.isnan(basins)]) + assert ids.size > 10 + for bid in ids: + members = basins == bid + # The outlet is the one cell in the basin that everything drains + # to, so it carries the basin's total. + outlet_total = float(np.max(accum[members])) + np.testing.assert_allclose( + outlet_total, float(np.sum(melt.values[members])), rtol=1e-9) + + +# --------------------------------------------------------------------------- +# 2. Watershed cross-check +# --------------------------------------------------------------------------- + +def _pour_points(fdir, accum, n_channel=8, n_random=8): + """Row/col pairs: the highest-accumulation cells plus random valid cells.""" + valid = ~np.isnan(fdir) + flat = np.where(valid, accum, -np.inf).ravel() + channel = np.argsort(flat)[::-1][:n_channel] + rng = np.random.default_rng(3734) + pool = np.flatnonzero(valid.ravel()) + random = rng.choice(pool, size=n_random, replace=False) + return [np.unravel_index(i, fdir.shape) for i in + np.concatenate([channel, random])] + + +@pytest.mark.parametrize('weight_name', ['melt', 'glacier']) +def test_watershed_cross_check(routers, melt, glacier, weight_name): + """accum[p] == sum(weight over watershed(p)) for many pour points.""" + fdir = routers['d8'] + weight = {'melt': melt, 'glacier': glacier}[weight_name] + accum = flow_accumulation(fdir, weight=weight).values + count = flow_accumulation(fdir).values + + for r, c in _pour_points(fdir.values, count): + pp = xr.full_like(fdir, np.nan) + pp.values[r, c] = 1.0 + ws = watershed(fdir, pp).values + inside = ws == 1.0 + assert inside[r, c] + np.testing.assert_allclose( + accum[r, c], float(np.sum(weight.values[inside])), rtol=1e-9) + # The unweighted count is the same identity with weight == 1. + assert count[r, c] == inside.sum() + + +def test_glacier_melt_is_zero_without_upstream_glacier(routers, glacier): + """Cells with no glacier anywhere upstream accumulate exactly zero.""" + fdir = routers['d8'] + accum = flow_accumulation(fdir, weight=glacier).values + basins = basin(fdir).values + for bid in np.unique(basins[~np.isnan(basins)]): + members = basins == bid + if glacier.values[members].sum() == 0: + assert np.all(accum[members] == 0.0) + else: + assert np.max(accum[members]) == glacier.values[members].sum() + + +# --------------------------------------------------------------------------- +# 3. Linearity +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize('routing', ['d8', 'dinf', 'mfd']) +def test_linearity(routers, melt, glacier, routing): + """accum(a*w1 + b*w2) == a*accum(w1) + b*accum(w2).""" + fdir = routers[routing] + a, b = 2.5, -0.75 + combined = a * melt + b * glacier + lhs = flow_accumulation(fdir, weight=combined, routing=routing).values + rhs = (a * flow_accumulation(fdir, weight=melt, routing=routing).values + + b * flow_accumulation(fdir, weight=glacier, routing=routing).values) + np.testing.assert_allclose(lhs, rhs, rtol=1e-9, atol=1e-9, equal_nan=True) + + +@pytest.mark.parametrize('routing', ['d8', 'dinf', 'mfd']) +def test_monotone_along_flow_paths(routers, melt, routing): + """With non-negative weight, accumulation never decreases downstream. + + For D8 this holds cell to cell. For D-inf and MFD each downstream + neighbour receives only a fraction, so check that the total handed + downstream (accum * fraction) is bounded by the receiver's value. + """ + fdir = routers[routing] + fd_np = fdir.values + accum = flow_accumulation(fdir, weight=melt, routing=routing).values + valid = ~np.isnan(fd_np[0] if routing == 'mfd' else fd_np) + h, w = valid.shape + checked = 0 + for r in range(h): + for c in range(w): + if not valid[r, c]: + continue + if routing == 'd8': + targets = [(_code_to_offset_py(fd_np[r, c]), 1.0)] + elif routing == 'dinf': + dy1, dx1, w1, dy2, dx2, w2 = _angle_to_neighbors(fd_np[r, c]) + targets = [((dy1, dx1), w1), ((dy2, dx2), w2)] + else: + dy = [0, 1, 1, 1, 0, -1, -1, -1] + dx = [1, 1, 0, -1, -1, -1, 0, 1] + targets = [((dy[k], dx[k]), fd_np[k, r, c]) for k in range(8)] + for (oy, ox), f in targets: + if f <= 0 or (oy, ox) == (0, 0): + continue + if _valid_at(valid, r + oy, c + ox): + assert accum[r + oy, c + ox] >= accum[r, c] * f - 1e-9 + checked += 1 + assert checked > 0