Repository navigation
weights for flow accumulation #3734
Description
Reason or Problem
flow_accumulation only computes a count of upstream cells draining through each cell. There's no way to weight that accumulation by another raster (e.g. precipitation or snowmelt), so getting a physically meaningful accumulated flux currently requires reimplementing the traversal outside the library.
Proposal
Add an optional weight parameter to flow_accumulation_d8 and the other flow_accumulation functions.
Design:
def flow_accumulation_d8(
flow_dir: xr.DataArray,
weight: xr.DataArray | None = None,
name: str = 'flow_accumulation',
) -> xr.DataArray:When weight is None, behavior is unchanged. When provided, each output cell equals the sum of weight over itself and all upstream cells draining through it, instead of a cell count. I.s. weight acts like a precipitation or surface melt field.
Usage:
melt_accum = flow_accumulation_d8(flow_dir, weight=melt)weight must be on the same grid as flow_dir.
Stakeholders and Impacts
Anyone using xarray-spatial for hydrology.
Unresolved Questions
How NaN values in weight should be handled (zero contribution vs. propagating as no-data downstream).
Activity
brendancol commentedon Aug 26, 2026
@jkingslake totally yes. Ive been doing some work on the surface_distance stuff so it makes sense to include a weight surface here. Really cool idea. I'd be happy to help review if you'd like to take a shot at it or if not, I'd be happy to address some time soon.
brendancol commentedon Aug 26, 2026
I think the answer to the "How to handle NaN" question is that they are zero contribution, like you have it in the melt example...so if there is a valid flow_dir cell, and a NaN weight, then the NaN would have no impact...also said, the output would NOT propagate a NaN value to the output based on that NaN weight.
I think we are on the same page here.
Address review: numpy-coerce cupy weights in MFD dask path, add dtype…
jkingslake commentedon Aug 26, 2026
Great! I am happy to have a go at this at some point. Probably with the help of claude.
I can see where currently uniform weights are assumed:
xarray-spatial/xrspatial/hydro/flow_accumulation_dinf.py
Line 215 in 037f569
xarray-spatial/xrspatial/hydro/flow_accumulation_dinf.py
Line 520 in 037f569
I guess that instead iterating through the new
weightsarray would be a good place to start:We would also need to take account of when
weights[r,c] = np.nanand makeaccum[r, c]zero.Would this be the main change do you think? (along with adding weights as an input etc.) Or would the be more involved things to do related to dask, cupy, cuda?
brendancol commentedon Aug 26, 2026
@jkingslake I tasked a claude agent with planning and the simple xarray-spatial loop, but we DON'T HAVE to go down that path if we don't want to.
From a library perspective, the hydrology tool spellings / api isn't perfect and I think it has some sprawl due to the d8 vs dinf, vs mdf...but let's ignore that for a moment.
The weight would need to be added to the datasets / datarray accessor methods and then each of the method specific / array-backends...basically a bunch of stamp collecting
brendancol commentedon Aug 26, 2026
@jkingslake I would love your help and don't mean to rush a PR up here. If you like to look at what the /rockout command from xarray-spatial-skills produced you can see it here: #3735
brendancol commentedon Aug 26, 2026
any chance you have example that we could run this against to test accuracy or build a test around?
Test weight through the Dataset accessor (#3734)
jkingslake commentedon Aug 26, 2026
I do not have anything to test this against particularly. The tests you have already looked at good and were along the lines of any synthetic tests I was imagining. I have a use case with an Antarctic DEM and fields of input from a regional climate model, but I don't think that will help test the implementation as such.
brendancol commentedon Aug 26, 2026
@jkingslake no worries. Maybe we could come up with a test that could confirm you or I's intuition about how it should work just on the simple create array example.
Add physical accuracy checks for weighted flow accumulation (#3734)
brendancol commentedon Aug 27, 2026
@jkingslake hey I've added some additional tests here that are more than just hand-rolled grids. I think we would probably gain more confidence through validating against pysheds...
some new tests are here: https://github.com/xarray-contrib/xarray-spatial/blob/990e5eddc74aad2e189fbe05f58bc5cf246fa81a/xrspatial/hydro/tests/test_flow_accumulation_weight_accuracy.py
Add weight raster to flow accumulation (#3734) (#3735)
brendancol commentedon Aug 27, 2026
This shipped in v0.10.18 (#3735), now on PyPI: https://github.com/xarray-contrib/xarray-spatial/releases/tag/v0.10.18
jkingslake commentedon Aug 27, 2026
Great work @brendancol !