Skip to content

Commit f08cb0f

Browse files
brendancolbrendancol
authored
Add native CUDA kernel for hydraulic erosion (#961) (#967)
* Add native CUDA kernel for hydraulic erosion (#961) Replaces the CPU-fallback path for CuPy and Dask+CuPy arrays with a real GPU kernel. Each particle maps to one CUDA thread; conflicts at shared heightmap cells are resolved with cuda.atomic.add. Refactors erode() to use ArrayTypeFunctionMapping dispatch instead of manual type checks. Also adds test_erosion.py with 15 tests covering numpy, dask+numpy, cupy, and dask+cupy backends. * Update README: hydraulic erosion now has native GPU support (#961) * Add hydraulic erosion user guide notebook (#961)
1 parent f4c9fcc commit f08cb0f

File tree

4 files changed

+649
-39
lines changed

4 files changed

+649
-39
lines changed

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,7 @@ In the GIS world, rasters are used for representing continuous phenomena (e.g. e
274274
| [Min Observable Height](xrspatial/experimental/min_observable_height.py) | Finds the minimum observer height needed to see each cell *(experimental)* | ✅️ | | | |
275275
| [Perlin Noise](xrspatial/perlin.py) | Generates smooth continuous random noise for procedural textures | ✅️ | ✅️ | ✅️ | ✅️ |
276276
| [Worley Noise](xrspatial/worley.py) | Generates cellular (Voronoi) noise returning distance to the nearest feature point | ✅️ | ✅️ | ✅️ | ✅️ |
277-
| [Hydraulic Erosion](xrspatial/erosion.py) | Simulates particle-based water erosion to carve valleys and deposit sediment | ✅️ | 🔄 | 🔄 | 🔄 |
277+
| [Hydraulic Erosion](xrspatial/erosion.py) | Simulates particle-based water erosion to carve valleys and deposit sediment | ✅️ | ✅️ | ✅️ | ✅️ |
278278
| [Bump Mapping](xrspatial/bump.py) | Adds randomized bump features to simulate natural terrain variation | ✅️ | ✅️ | ✅️ | ✅️ |
279279

280280
-----------

‎examples/user_guide/21_Hydraulic_Erosion.ipynb‎

Lines changed: 235 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,235 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "cell-0",
6+
"metadata": {},
7+
"source": [
8+
"# Hydraulic Erosion\n",
9+
"\n",
10+
"The `erode()` function simulates particle-based hydraulic erosion on a terrain raster. Thousands of virtual water droplets are dropped on the surface, each one tracing a path downhill. Along the way, droplets pick up sediment from steep slopes and deposit it when they slow down or encounter flatter ground.\n",
11+
"\n",
12+
"This produces realistic valley networks and smoothed ridgelines from synthetic or real terrain data. The function supports numpy, cupy, dask+numpy, and dask+cupy backends. On GPU, each particle runs as a separate CUDA thread for large speedups on high-resolution grids."
13+
]
14+
},
15+
{
16+
"cell_type": "code",
17+
"execution_count": null,
18+
"id": "cell-1",
19+
"metadata": {},
20+
"outputs": [],
21+
"source": [
22+
"import numpy as np\n",
23+
"import xarray as xr\n",
24+
"import matplotlib.pyplot as plt\n",
25+
"\n",
26+
"from xrspatial.erosion import erode"
27+
]
28+
},
29+
{
30+
"cell_type": "markdown",
31+
"id": "cell-2",
32+
"metadata": {},
33+
"source": [
34+
"## 1. Generate synthetic terrain\n",
35+
"\n",
36+
"We build a simple terrain using layered sine waves plus random noise. This gives enough slope variation for erosion to produce visible channels."
37+
]
38+
},
39+
{
40+
"cell_type": "code",
41+
"execution_count": null,
42+
"id": "cell-3",
43+
"metadata": {},
44+
"outputs": [],
45+
"source": [
46+
"size = 256\n",
47+
"y = np.linspace(0, 4 * np.pi, size)\n",
48+
"x = np.linspace(0, 4 * np.pi, size)\n",
49+
"xx, yy = np.meshgrid(x, y)\n",
50+
"\n",
51+
"# Layered sine terrain with noise\n",
52+
"rng = np.random.default_rng(42)\n",
53+
"terrain = (\n",
54+
" 200 * np.sin(xx * 0.3) * np.cos(yy * 0.2)\n",
55+
" + 100 * np.sin(xx * 0.7 + yy * 0.5)\n",
56+
" + 50 * rng.random((size, size))\n",
57+
" + 500 # raise the baseline so values stay positive\n",
58+
")\n",
59+
"\n",
60+
"agg = xr.DataArray(\n",
61+
" terrain.astype(np.float32),\n",
62+
" dims=['y', 'x'],\n",
63+
" attrs={'res': (1.0, 1.0)},\n",
64+
")\n",
65+
"\n",
66+
"fig, ax = plt.subplots(figsize=(8, 6))\n",
67+
"im = ax.imshow(agg.values, cmap='terrain')\n",
68+
"ax.set_title('Original terrain')\n",
69+
"ax.axis('off')\n",
70+
"fig.colorbar(im, ax=ax, shrink=0.7, label='Elevation')\n",
71+
"plt.tight_layout()\n",
72+
"plt.show()"
73+
]
74+
},
75+
{
76+
"cell_type": "markdown",
77+
"id": "cell-4",
78+
"metadata": {},
79+
"source": [
80+
"## 2. Basic erosion\n",
81+
"\n",
82+
"Run erosion with default parameters and compare against the original."
83+
]
84+
},
85+
{
86+
"cell_type": "code",
87+
"execution_count": null,
88+
"id": "cell-5",
89+
"metadata": {},
90+
"outputs": [],
91+
"source": [
92+
"eroded = erode(agg, iterations=50000, seed=42)\n",
93+
"\n",
94+
"fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n",
95+
"\n",
96+
"axes[0].imshow(agg.values, cmap='terrain')\n",
97+
"axes[0].set_title('Before')\n",
98+
"axes[0].axis('off')\n",
99+
"\n",
100+
"axes[1].imshow(eroded.values, cmap='terrain')\n",
101+
"axes[1].set_title('After erosion')\n",
102+
"axes[1].axis('off')\n",
103+
"\n",
104+
"diff = eroded.values - agg.values\n",
105+
"vlim = max(abs(diff.min()), abs(diff.max()))\n",
106+
"im = axes[2].imshow(diff, cmap='RdBu', vmin=-vlim, vmax=vlim)\n",
107+
"axes[2].set_title('Elevation change')\n",
108+
"axes[2].axis('off')\n",
109+
"fig.colorbar(im, ax=axes[2], shrink=0.7, label='Change')\n",
110+
"\n",
111+
"plt.tight_layout()\n",
112+
"plt.show()"
113+
]
114+
},
115+
{
116+
"cell_type": "markdown",
117+
"id": "cell-6",
118+
"metadata": {},
119+
"source": [
120+
"## 3. Parameter effects\n",
121+
"\n",
122+
"The `params` dict controls the erosion behavior. The most impactful ones:\n",
123+
"\n",
124+
"- **erosion**: how aggressively particles remove material (default 0.3)\n",
125+
"- **capacity**: how much sediment water can carry (default 4.0)\n",
126+
"- **deposition**: how quickly excess sediment is dropped (default 0.3)\n",
127+
"- **radius**: size of the erosion brush in cells (default 3)"
128+
]
129+
},
130+
{
131+
"cell_type": "code",
132+
"execution_count": null,
133+
"id": "cell-7",
134+
"metadata": {},
135+
"outputs": [],
136+
"source": [
137+
"configs = {\n",
138+
" 'Default': None,\n",
139+
" 'High erosion': {'erosion': 0.9},\n",
140+
" 'Large brush': {'radius': 6},\n",
141+
" 'High capacity': {'capacity': 12.0},\n",
142+
"}\n",
143+
"\n",
144+
"fig, axes = plt.subplots(1, 4, figsize=(20, 5))\n",
145+
"for ax, (label, p) in zip(axes, configs.items()):\n",
146+
" result = erode(agg, iterations=30000, seed=42, params=p)\n",
147+
" diff = result.values - agg.values\n",
148+
" vlim = max(abs(diff.min()), abs(diff.max()), 1)\n",
149+
" ax.imshow(diff, cmap='RdBu', vmin=-vlim, vmax=vlim)\n",
150+
" ax.set_title(label)\n",
151+
" ax.axis('off')\n",
152+
"\n",
153+
"plt.suptitle('Elevation change under different parameters', y=1.02)\n",
154+
"plt.tight_layout()\n",
155+
"plt.show()"
156+
]
157+
},
158+
{
159+
"cell_type": "markdown",
160+
"id": "cell-8",
161+
"metadata": {},
162+
"source": [
163+
"## 4. Iterative refinement\n",
164+
"\n",
165+
"More iterations means more droplets, which produces deeper, more developed channel networks."
166+
]
167+
},
168+
{
169+
"cell_type": "code",
170+
"execution_count": null,
171+
"id": "cell-9",
172+
"metadata": {},
173+
"outputs": [],
174+
"source": [
175+
"iter_counts = [5000, 20000, 50000, 100000]\n",
176+
"\n",
177+
"fig, axes = plt.subplots(1, 4, figsize=(20, 5))\n",
178+
"for ax, n in zip(axes, iter_counts):\n",
179+
" result = erode(agg, iterations=n, seed=42)\n",
180+
" ax.imshow(result.values, cmap='terrain')\n",
181+
" ax.set_title(f'{n:,} droplets')\n",
182+
" ax.axis('off')\n",
183+
"\n",
184+
"plt.suptitle('Erosion depth vs. iteration count', y=1.02)\n",
185+
"plt.tight_layout()\n",
186+
"plt.show()"
187+
]
188+
},
189+
{
190+
"cell_type": "markdown",
191+
"id": "cell-10",
192+
"metadata": {},
193+
"source": [
194+
"## 5. Cross-section comparison\n",
195+
"\n",
196+
"A 1D slice through the terrain shows how erosion carves valleys and smooths peaks."
197+
]
198+
},
199+
{
200+
"cell_type": "code",
201+
"execution_count": null,
202+
"id": "cell-11",
203+
"metadata": {},
204+
"outputs": [],
205+
"source": [
206+
"row = size // 2\n",
207+
"eroded_heavy = erode(agg, iterations=80000, seed=42)\n",
208+
"\n",
209+
"fig, ax = plt.subplots(figsize=(12, 4))\n",
210+
"ax.plot(agg.values[row, :], label='Original', linewidth=2)\n",
211+
"ax.plot(eroded.values[row, :], label='50k droplets', linewidth=1.5)\n",
212+
"ax.plot(eroded_heavy.values[row, :], label='80k droplets', linewidth=1.5)\n",
213+
"ax.set_xlabel('Column')\n",
214+
"ax.set_ylabel('Elevation')\n",
215+
"ax.set_title(f'Cross-section at row {row}')\n",
216+
"ax.legend()\n",
217+
"plt.tight_layout()\n",
218+
"plt.show()"
219+
]
220+
}
221+
],
222+
"metadata": {
223+
"kernelspec": {
224+
"display_name": "Python 3",
225+
"language": "python",
226+
"name": "python3"
227+
},
228+
"language_info": {
229+
"name": "python",
230+
"version": "3.10.0"
231+
}
232+
},
233+
"nbformat": 4,
234+
"nbformat_minor": 5
235+
}

0 commit comments

Comments
 (0)