Add weight raster to flow accumulation (#3734)

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",