Add weight raster to flow accumulation (#3734)

[brendancol]

Closes #3734

Adds an optional weight raster to flow_accumulation_d8, flow_accumulation_dinf and flow_accumulation_mfd. flow_accumulation(routing=...) and the .xrs accessor forward kwargs, so they pick it up too.

With weight, each output cell is the sum of weight over itself and everything upstream of it (split by the D-inf / MFD fractions for those routers). Without it nothing changes: same unit count, same kernels, same memory footprint. The kernels take (weight, has_weight) and the unweighted path passes a 1x1 dummy.

NaN handling, answering the open question on the issue: a NaN weight contributes 0 and the output NaN mask still follows flow_dir. Propagating NaN downstream would blank the whole basin below one missing pixel, and treating it as a barrier would silently drop upstream flow, so neither seemed usable. Anyone wanting a different convention can fillna first.

Other details:

• weight sits between flow_dir and name, as proposed on the issue, so name is no longer the second positional argument. Nothing in the repo passes name positionally (accessor and notebooks use the keyword) and a stray string now raises a clear TypeError from _validate_raster.

• weight goes through _validate_raster + _validate_matching_shape; a shape mismatch raises ValueError.

• Memory guards add 8 bytes/pixel when a weight is present.

• Dask paths carry the weight tile through the boundary sweep and the lazy map_blocks assembly. A numpy weight, or a dask weight chunked differently from flow_dir, is rechunked to match.

• Docstrings updated, plus a short "Weighted accumulation" section in the hydrology user guide notebook.

Backends: numpy, cupy, dask+numpy, dask+cupy for all three routers.

Test plan:

• Known-value tests per router (D8 chain, D-inf 8-neighbour pit and a pi/8 proportional split, MFD plane)
• weight=1 reproduces the unweighted count
• NaN weight contributes 0; NaN mask equals isnan(flow_dir)
• Shape mismatch and non-DataArray weight raise
• dask parity across chunk layouts, including a weight chunked differently from flow_dir
• cupy and dask+cupy parity (run locally on a CUDA box)
• Physical accuracy checks on a generate_terrain DEM with an orographic melt field (test_flow_accumulation_weight_accuracy.py): mass balance per router and per basin (input weight = pit totals + leak off the valid grid, leak computed from flow_dir independently of the kernels), watershed cross-check at channel and hillslope pour points with melt and a glacier indicator, linearity in weight, and downstream monotonicity for non-negative weight. Mass balance runs on all four backends.
• Existing xrspatial/hydro/tests and accessor tests still pass (1083 passed)

[brendancol]

PR Review: Add weight raster to flow accumulation (#3734)

Verified in a worktree on the PR head. The three test files pass (121 passed, including the cupy and dask+cupy cases on a CUDA box). Also ran ad-hoc probes for every backend pairing of flow_dir and weight (numpy, cupy, dask+numpy, dask+cupy on both sides), int/inf/3-D weights, Dataset input, routing= dispatch and positional name.

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

• xrspatial/hydro/flow_accumulation_mfd.py:714-716 and :836: the MFD dask+numpy path converts the weight tile with a bare np.asarray(...), so a cupy-backed weight with a dask+numpy flow_dir_mfd raises "Implicit conversion to a NumPy array is not allowed" (reproduced). The D8 and D-inf paths route the same tile through _to_numpy_f64, which handles .get(), and every other mixed pairing works (cupy fd + numpy w, dask+cupy fd + numpy w, dask+numpy fd + cupy w for D8/D-inf, and all MFD pairings except this one). Either use _to_numpy_f64 here too, or reject mismatched backends up front in _validate_weight. Low priority since it is a mixed-backend corner, but it is the one hole in the parity matrix.
• xrspatial/hydro/flow_accumulation_d8.py:1071-1073 (and dinf :985-987, mfd :879-881): weight is inserted before name, so flow_accumulation_d8(fd, 'myname') now raises TypeError: weight must be an xarray.DataArray, got builtins.str instead of naming the output. No positional name callers in the repo (accessor and notebooks use kwargs), and the error message is clear, so this is a judgment call. If you want to keep the old positional contract, append weight after name or make both keyword-only. Worth a line in the PR description either way.
• Tests do not cover: an integer-dtype weight (works via np.asarray(dtype=float64)), a Dataset flow_dir with a DataArray weight through @supports_dataset (works, each variable gets the same weight), and the .xrs accessor / flow_accumulation(routing=...) forwarding path the PR description relies on (works). One small parametrised test for the first two would lock in the coercion behaviour.

Nits (optional improvements)

• xrspatial/hydro/tests/test_flow_accumulation_dinf.py:321: docstring says "flowing at pi/4 is split 50/50" but the case uses pi/8 (the comment on the next line is correct). Fix the docstring.
• xrspatial/hydro/tests/test_flow_accumulation_mfd.py:530-541: test_mass_conserved_on_plane docstring describes a bottom-row column-sum check, but the assertion is accum == count * 2.5 (uniform-weight scaling). Either rename or make the docstring match.
• xrspatial/hydro/tests/test_flow_accumulation_dinf.py:365-368: every parametrised chunking uses a weight chunked (2, 7), so the "weight already chunked like flow_dir" branch of _weight_as_dask is only exercised by the D8 and MFD tests. Fine as-is since the helper is shared.
• xrspatial/hydro/flow_accumulation_d8.py:806-810: the weight tile is recomputed from the dask graph on every sweep pass. This mirrors how flow_dir_da.blocks[iy, ix].compute() is already handled, so it is not a regression, but a weight backed by an expensive graph (e.g. a lazy reprojection) pays that cost twice per sweep. Not worth changing here.

What looks good

• Kernel indexing is correct everywhere checked: _cell_weight (d8.py:154) and the three @cuda.jit init kernels index weight[r, c] / weight[i, j] with the same local tile coordinates used for flow_dir, and _NO_WEIGHT is never indexed when has_weight == 0.
• NaN policy is implemented identically on all four backends for all three routers: NaN weight becomes 0.0, NaN mask stays isnan(flow_dir) (or isnan(fractions[0]) for MFD). Negative and inf weights propagate arithmetically, no clamping, matching the documented contract.
• Dask tile alignment: _weight_as_dask (d8.py:177) rechunks to flow_dir.chunks before any .blocks[iy, ix] access, so the D8/D-inf tile lookups and the MFD cumulative-offset slicing (mfd.py:697-716) always see the same tile. Verified with (2,5) / (4,2) weight chunks against (3,3) flow chunks and with neighbor chunked on axis 0.
• map_blocks with two inputs: block_info[0] still refers to the flow array in both _assemble_result* variants, and the 2-D weight alongside the 3-D fractions with drop_axis=0 (mfd.py:840-846) pairs correctly by trailing axes. dask+cupy with a numpy weight still yields cupy blocks.
• No eager materialisation: a numpy weight is wrapped with da.from_array, a dask weight stays lazy, tiles are computed one at a time.
• Memory guards: 8 bytes/pixel for the float64 weight copy is right for the conversion cases and slightly conservative when the weight is already float64, which is the safe direction.
• Notebook cells come after dem, H, W, flow_dir, flow_accum, water_cmap and LogNorm are defined; the LogNorm vmin guard is fine because precip is strictly positive.
• Docstrings are accurate for all three routers.

Checklist

• Algorithm matches reference (sum of weight over self + upstream; fraction-split for D-inf/MFD)
• Backends consistent (one mixed-backend corner in MFD noted above)
• NaN handling
• Edge cases tested (int dtype and Dataset input untested)
• Dask chunk boundaries
• No premature materialisation
• Benchmark exists or not needed (weighted path only changes per-cell init)
• README updated if applicable (n/a)
• Docstrings

[brendancol]

PR Review: follow-up after 1d802e0

Re-checked the review findings against the follow-up commit.

Blockers (must fix before merge)

None.

Suggestions (should fix, not blocking)

None remaining.

Disposition of the first-pass findings

• Fixed: MFD dask+numpy path with a cupy weight. xrspatial/hydro/flow_accumulation_mfd.py:715, :836 and :953 now go through _to_numpy_f64, which calls .get() on cupy input and is a no-op on numpy. Reproduced the failure before the change and confirmed it passes after.
• Fixed: added test_weight_integer_dtype, test_weight_dataset_input and test_weight_forwarded_by_accessor_and_routing in test_flow_accumulation_d8.py. 124 tests pass across the three files.
• Fixed: both test docstrings (test_flow_accumulation_dinf.py pi/8 wording; MFD test renamed to test_uniform_weight_scales_count with a matching docstring), and the D-inf dask test now also exercises a weight chunked identically to flow_dir.
• Dismissed with reason: positional name ordering. Kept weight before name to match the signature proposed on #3734; no in-repo positional callers and the failure mode is a clear TypeError. Noted in the PR description.
• Dismissed with reason: weight tile recomputed per sweep pass. Mirrors the existing flow_dir handling; reviewer agreed it is not worth changing here.

Checklist

• Backends consistent (mixed-backend MFD corner now covered)
• Tests pass locally including cupy and dask+cupy

[brendancol]

Ad-hoc cross-check against pysheds 0.5 (not part of the PR, run locally).

Setup: 150x150 generate_terrain DEM (seed 3734), flow_direction_d8 / flow_direction_dinf from xrspatial, and an orographic melt field as the weight with a 4x4 NaN patch. Both libraries got the identical flow_dir; pysheds got the NaN patch as explicit zeros to match the NaN-contributes-0 policy. Direction codes map as xrspatial E=1 SE=2 S=4 SW=8 W=16 NW=32 N=64 NE=128 to pysheds dirmap=(64, 128, 1, 2, 4, 8, 16, 32).

router	valid cells	max abs diff	max rel diff	outlet total (xrspatial / pysheds)
D8 weighted	21,904	4.5e-13	5.4e-16	4628.2563 / 4628.2563
D8 count (unweighted sanity)	21,904	0	0	1335 / 1335
D-inf weighted	21,904	9.1e-13	8.3e-16	4608.0347 / 4608.0347

That is float64 rounding on every valid cell for both routers, NaN-weight cells included. The exact match on the unweighted count confirms the code mapping, so the weighted agreement is not an artefact of a misaligned grid.

MFD is not compared: pysheds partitions MFD flow with a different scheme than flow_direction_mfd, so the fractions differ before accumulation starts. For MFD the evidence is the mass-balance and linearity tests in test_flow_accumulation_weight_accuracy.py.

If it is useful as a permanent test it fits the reference-validation pattern with pysheds as an optional dependency; happy to do that as a follow-up.