Skip to content

Add hydrological flow visualization #70

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
brendancol wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Add hydrological flow visualization #70

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0cc0733
Add hydrological flow visualization and update Trinidad example
brendancol Mar 1, 2026
345841f
Improve hydro flow: waterway DEM burning, water shader overlay, polyg…
brendancol Mar 2, 2026
7662d66
Add tributary flow tracing and coast distance analysis to quickstart
brendancol Mar 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 83 additions & 1 deletion examples/playground.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,88 @@ def load_terrain():
except Exception as e:
print(f"Skipping wind: {e}")

print("\nLaunching explore (press G to cycle layers, Shift+W for wind)...\n")
# --- Hydro flow data (off by default, Shift+Y to toggle) ---------------
hydro = None
try:
from xrspatial import fill as _fill
from xrspatial import flow_direction as _flow_direction
from xrspatial import flow_accumulation as _flow_accumulation
from xrspatial import stream_order as _stream_order
from xrspatial import stream_link as _stream_link
from scipy.ndimage import uniform_filter as _uniform_filter

print("Conditioning DEM for hydrological flow...")
_data = terrain.data
_is_cupy = hasattr(_data, 'get')
_elev = _data.get() if _is_cupy else np.array(_data)
_elev = _elev.astype(np.float32)
_ocean = (_elev == 0.0) | np.isnan(_elev)
_elev[_ocean] = -100.0

_smoothed = _uniform_filter(_elev, size=15, mode='nearest')
_smoothed[_ocean] = -100.0

if _is_cupy:
import cupy as _cp
_sm = _cp.asarray(_smoothed)
else:
_sm = _smoothed
_filled = _fill(terrain.copy(data=_sm))
_fd = _filled.data - _sm
_resolved = _filled.data + _fd * 0.01
if _is_cupy:
_cp.random.seed(0)
_resolved += _cp.random.uniform(0, 0.001, _resolved.shape,
dtype=_cp.float32)
_resolved[_cp.asarray(_ocean)] = -100.0
else:
np.random.seed(0)
_resolved += np.random.uniform(
0, 0.001, _resolved.shape).astype(np.float32)
_resolved[_ocean] = -100.0

fd = _flow_direction(terrain.copy(data=_resolved))
fa = _flow_accumulation(fd)
so = _stream_order(fd, fa, threshold=50)
sl = _stream_link(fd, fa, threshold=50)

fd_out, fa_out, so_out = fd.data, fa.data, so.data
if _is_cupy:
_ocean_gpu = _cp.asarray(_ocean)
fd_out[_ocean_gpu] = _cp.nan
fa_out[_ocean_gpu] = _cp.nan
so_out[_ocean_gpu] = _cp.nan
else:
fd_out[_ocean] = np.nan
fa_out[_ocean] = np.nan
so_out[_ocean] = np.nan

_sl_out = sl.data
if _is_cupy:
_sl_out[_ocean_gpu] = _cp.nan
else:
_sl_out[_ocean] = np.nan
_sl_np = _sl_out.get() if _is_cupy else np.asarray(_sl_out)
_sl_clean = np.nan_to_num(_sl_np, nan=0.0).astype(np.float32)
if _is_cupy:
_sl_clean = _cp.asarray(_sl_clean)
ds['stream_link'] = terrain.copy(data=_sl_clean).rename(None)

hydro = {
'flow_dir': fd_out,
'flow_accum': fa_out,
'stream_order': so_out,
'stream_link': _sl_out,
'accum_threshold': 50,
'enabled': False,
}
print(f" Flow direction + accumulation computed on "
f"{terrain.shape[0]}x{terrain.shape[1]} grid")
except Exception as e:
print(f"Skipping hydro: {e}")

print("\nLaunching explore (press G to cycle layers, "
"Shift+W for wind, Shift+Y for hydro)...\n")
ds.rtx.explore(
z='elevation',
scene_zarr=ZARR,
Expand All @@ -245,6 +326,7 @@ def load_terrain():
height=768,
render_scale=0.5,
wind_data=wind,
hydro_data=hydro,
repl=True,
)

Expand Down
4 changes: 3 additions & 1 deletion examples/trinidad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""Trinidad and Tobago — GPU-accelerated terrain exploration."""
"""Trinidad and Tobago — GPU-accelerated terrain exploration with hydro flow."""
from rtxpy import quickstart

quickstart(
Expand All @@ -7,6 +7,8 @@
crs='EPSG:32620',
features=['buildings', 'roads', 'water', 'fire', 'places',
'infrastructure', 'land_use'],
hydro=True,
coast_distance=True,
ao_samples=1,
denoise=True,
)
7 changes: 5 additions & 2 deletions rtxpy/accessor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2547,7 +2547,7 @@ def explore(self, width=800, height=600, render_scale=0.5,
start_position=None, look_at=None, key_repeat_interval=0.05,
pixel_spacing_x=None, pixel_spacing_y=None,
mesh_type='heightfield', color_stretch='linear', title=None,
subsample=1, wind_data=None, gtfs_data=None,
subsample=1, wind_data=None, hydro_data=None, gtfs_data=None,
terrain_loader=None,
scene_zarr=None, ao_samples=0, gi_bounces=1, denoise=False,
repl=False, tour=None):
Expand Down Expand Up @@ -2686,6 +2686,7 @@ def explore(self, width=800, height=600, render_scale=0.5,
baked_meshes=self._baked_meshes if self._baked_meshes else None,
subsample=subsample,
wind_data=wind_data,
hydro_data=hydro_data,
gtfs_data=gtfs_data,
accessor=self,
terrain_loader=terrain_loader,
Expand Down Expand Up @@ -3007,7 +3008,8 @@ def explore(self, z, width=800, height=600, render_scale=0.5,
pixel_spacing_x=None, pixel_spacing_y=None,
mesh_type='heightfield', color_stretch='linear', title=None,
subtitle=None, legend=None,
subsample=1, wind_data=None, gtfs_data=None,
subsample=1, wind_data=None, hydro_data=None,
gtfs_data=None,
scene_zarr=None,
ao_samples=0, gi_bounces=1, denoise=False,
minimap_style=None, minimap_layer=None,
Expand Down Expand Up @@ -3113,6 +3115,7 @@ def explore(self, z, width=800, height=600, render_scale=0.5,
baked_meshes=terrain_da.rtx._baked_meshes if terrain_da.rtx._baked_meshes else None,
subsample=subsample,
wind_data=wind_data,
hydro_data=hydro_data,
gtfs_data=gtfs_data,
accessor=terrain_da.rtx,
scene_zarr=scene_zarr,
Expand Down
43 changes: 31 additions & 12 deletions rtxpy/analysis/render.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -823,7 +823,7 @@ def _shade_terrain_kernel(
color_stretch,
rgb_texture,
overlay_data, overlay_alpha, overlay_min, overlay_range,
overlay_as_water,
overlay_as_water, overlay_color_lut,
instance_ids, geometry_colors,
primitive_ids, point_colors, point_color_offsets,
ao_factor, gi_color, gi_intensity,
Expand Down Expand Up @@ -1006,21 +1006,28 @@ def _shade_terrain_kernel(
ov_norm = 0.0
elif ov_norm > 1:
ov_norm = 1.0
# Apply same color stretch
if color_stretch == 1:
ov_norm = math.pow(ov_norm, 1.0 / 3.0)
elif color_stretch == 2:
ov_norm = math.log(1.0 + ov_norm * 9.0) / math.log(10.0)
elif color_stretch == 3:
ov_norm = math.sqrt(ov_norm)
_use_ov_lut = overlay_color_lut.shape[0] > 1
if not _use_ov_lut:
# Apply color stretch only with terrain colormap
if color_stretch == 1:
ov_norm = math.pow(ov_norm, 1.0 / 3.0)
elif color_stretch == 2:
ov_norm = math.log(1.0 + ov_norm * 9.0) / math.log(10.0)
elif color_stretch == 3:
ov_norm = math.sqrt(ov_norm)
ov_idx = int(ov_norm * 255)
if ov_idx > 255:
ov_idx = 255
if ov_idx < 0:
ov_idx = 0
ov_r = color_lut[ov_idx, 0]
ov_g = color_lut[ov_idx, 1]
ov_b = color_lut[ov_idx, 2]
if _use_ov_lut:
ov_r = overlay_color_lut[ov_idx, 0]
ov_g = overlay_color_lut[ov_idx, 1]
ov_b = overlay_color_lut[ov_idx, 2]
else:
ov_r = color_lut[ov_idx, 0]
ov_g = color_lut[ov_idx, 1]
ov_b = color_lut[ov_idx, 2]
a = overlay_alpha
base_r = base_r * (1.0 - a) + ov_r * a
base_g = base_g * (1.0 - a) + ov_g * a
Expand Down Expand Up @@ -1552,6 +1559,7 @@ def _shade_terrain(
overlay_data=None, overlay_alpha=0.5,
overlay_min=0.0, overlay_range=1.0,
overlay_as_water=False,
overlay_color_lut=None,
instance_ids=None, geometry_colors=None,
primitive_ids=None, point_colors=None, point_color_offsets=None,
ao_factor=None, gi_color=None, gi_intensity=2.0,
Expand Down Expand Up @@ -1585,6 +1593,15 @@ def _shade_terrain(
_DUMMY_1x1 = cupy.zeros((1, 1), dtype=np.float32)
overlay_data = _DUMMY_1x1

# Handle overlay color LUT — custom palette for categorical overlays
# Dummy is (1,3); kernel checks shape[0] > 1 to decide whether to use
if overlay_color_lut is None:
if _DUMMY_1x1x3 is None:
_DUMMY_1x1x3 = cupy.zeros((1, 1, 3), dtype=np.float32)
overlay_color_lut = _DUMMY_1x1x3[:, 0, :] # (1, 3)
elif not isinstance(overlay_color_lut, cupy.ndarray):
overlay_color_lut = cupy.asarray(overlay_color_lut, dtype=np.float32)

# Handle geometry_colors for per-geometry solid coloring
if geometry_colors is None:
if _DUMMY_1x4 is None:
Expand Down Expand Up @@ -1655,7 +1672,7 @@ def _shade_terrain(
color_stretch,
rgb_texture,
overlay_data, overlay_alpha, overlay_min, overlay_range,
overlay_as_water,
overlay_as_water, overlay_color_lut,
instance_ids, geometry_colors,
primitive_ids, point_colors, point_color_offsets,
ao_factor, gi_color, np.float32(gi_intensity),
Expand Down Expand Up @@ -1788,6 +1805,7 @@ def render(
overlay_data=None,
overlay_alpha: float = 0.5,
overlay_as_water: bool = False,
overlay_color_lut=None,
geometry_colors=None,
ao_samples: int = 0,
ao_radius: Optional[float] = None,
Expand Down Expand Up @@ -2166,6 +2184,7 @@ def render(
overlay_data=d_overlay, overlay_alpha=overlay_alpha,
overlay_min=ov_min, overlay_range=ov_range,
overlay_as_water=overlay_as_water,
overlay_color_lut=overlay_color_lut,
instance_ids=d_instance_ids, geometry_colors=geometry_colors,
primitive_ids=d_primitive_ids,
point_colors=d_point_colors,
Expand Down
Loading
Loading