weights for flow accumulation

[jkingslake]

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).

[brendancol]

@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]

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.

[jkingslake]

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

accum[r, c] = 1.0

and

xarray-spatial/xrspatial/hydro/flow_accumulation_dinf.py

Line 520 in 037f569

accum[r, c] = 1.0

I guess that instead iterating through the new weights array would be a good place to start:

accum[r, c] = weights[r, c]

We would also need to take account of when weights[r,c] = np.nan and make accum[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]

@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]

@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]

any chance you have example that we could run this against to test accuracy or build a test around?

[jkingslake]

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]

@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.