diff --git a/.claude/commands/rockout.md b/.claude/commands/rockout.md new file mode 100644 index 0000000..76ebd51 --- /dev/null +++ b/.claude/commands/rockout.md @@ -0,0 +1,107 @@ +# Rockout: End-to-End Issue-to-Implementation Workflow + +Take the user's prompt describing an enhancement, bug, or suggestion and drive it +through all seven steps below. The prompt is: $ARGUMENTS + +--- + +## Step 1 -- Create a GitHub Issue + +1. Decide the issue type from the prompt: + - **enhancement** -- new feature or improvement + - **bug** -- something broken + - **suggestion / proposal** -- idea that needs design discussion +2. Pick labels from the repo's existing set. Always include the type label + (`enhancement`, `bug`). Add topical labels when they fit + (e.g. `documentation`, `GPU CI`). +3. Draft the title and body. Write a clear problem statement, motivation, and + proposed approach. For bugs, include reproduction steps and expected vs actual + behavior. +4. **Run the body text through the `/humanizer` skill** before creating the issue + to strip AI writing patterns. +5. Create the issue with `gh issue create` using the drafted title, body, and labels. +6. Capture the new issue number for later steps. + +## Step 2 -- Create a Git Worktree + +1. Create a new branch and worktree using the issue number: + ``` + git worktree add .claude/worktrees/issue- -b issue- + ``` +2. Switch the working directory to the new worktree for all remaining steps. + +## Step 3 -- Implement the Change + +1. Read the relevant source files to understand the existing code. +2. Follow the project's architecture patterns: + - **xarray accessor**: public API goes through `rtxpy/accessor.py` (`RTXAccessor`) + - **Engine/viewer**: interactive features in `rtxpy/engine.py` and `rtxpy/viewer/` subsystems + - **Analysis**: raster operations in `rtxpy/analysis/` with `prepare_mesh()` in `_common.py` + - **Mesh utils**: geometry operations in `rtxpy/mesh.py` + - **Data fetching**: `fetch_*` functions for remote data sources +3. When adding viewer features, use the key dispatch pattern: + add entries to `KEY_BINDINGS` / `SHIFT_BINDINGS` / `SPECIAL_BINDINGS` in + `rtxpy/viewer/keybindings.py`, then add thin action methods on `InteractiveViewer`. +4. When adding analysis functions, return standard xarray DataArrays. +5. Keep changes focused -- don't refactor surrounding code unnecessarily. + +## Step 4 -- Add Test Coverage + +1. Add or update tests in `rtxpy/tests/`. +2. Any temporary files must have unique names. Include the issue number in + the filename (e.g. `tmp_42_result.zarr`) to avoid collisions with + parallel test runs or other worktrees. +3. Cover: + - Correctness against known values or reference implementations + - Edge cases (NaN handling, empty input, single-cell rasters) + - GPU and CPU code paths where applicable +4. Run the tests with `pytest rtxpy/tests/` to verify they pass before moving on. + +## Step 5 -- Update Documentation + +1. Check `docs/` for the relevant markdown file: + - `docs/api-reference.md` -- method signatures and parameters + - `docs/user-guide.md` -- task-oriented workflows + - `docs/getting-started.md` -- installation and first steps + - `docs/examples.md` -- annotated walkthroughs +2. Add or update the entry for any new public functions or viewer features. +3. If a new data fetcher or analysis function was created, add it to the + appropriate section. + +## Step 6 -- Create an Example + +The project has an `examples/` directory with both `.py` scripts and `.ipynb` notebooks. + +1. Choose the right format: + - `.py` script for CLI-oriented workflows or `explore()` demos + - `.ipynb` notebook for analysis workflows with inline visualization +2. Follow the established patterns from existing examples: + - Fetch data from a bounding box using `fetch_*` functions + - Show the analysis or placement pipeline + - End with `explore()` or a rendered image +3. Use `matplotlib` for static plots, consistent with existing notebooks. +4. Keep the example self-contained (no external data dependencies beyond fetch calls). + +**Skip this step** if the change is a pure bug fix with no new user-facing API. + +## Step 7 -- Commit, Push, and Create PR + +1. Stage and commit changes with a clear message referencing the issue number + (e.g. `Add flood velocity function (#42)`). +2. Push the branch and create a pull request with `gh pr create`. +3. In the PR body, reference the issue (e.g. `Closes #42`) and summarize + what was done. +4. **Run the PR body text through the `/humanizer` skill** before creating. + +--- + +## General Rules + +- Work entirely within the worktree created in Step 2. +- Commit progress after each major step with a clear commit message referencing + the issue number. +- Run `/humanizer` on any text destined for GitHub (issue body, PR description) + to remove AI writing artifacts. +- If any step is not applicable (e.g. no docs update needed for a typo fix), + note why and skip it. +- At the end, print a summary of what was done and where the worktree lives. diff --git a/examples/bay_area_fire_risk.py b/examples/bay_area_fire_risk.py index c33b603..a01bf42 100644 --- a/examples/bay_area_fire_risk.py +++ b/examples/bay_area_fire_risk.py @@ -513,6 +513,14 @@ def cluster_buildings(bldg_data, terrain_da): except Exception as e: print(f" Skipping wind: {e}") + # === Fetch weather data (clouds + rain via Shift+N) ======================= + weather = None + try: + from rtxpy import fetch_weather + weather = fetch_weather(BOUNDS, grid_size=15) + except Exception as e: + print(f" Skipping weather: {e}") + print("\n--- Interpolating wind onto DEM grid ---") wind_speed_gpu, wind_u_gpu, wind_v_gpu = interpolate_wind(wind, terrain) ds['wind_speed'] = xr.DataArray( @@ -695,6 +703,7 @@ def cluster_buildings(bldg_data, terrain_da): color_stretch='cbrt', subsample=1, wind_data=wind, + weather_data=weather, ao_samples=1, repl=True, ) diff --git a/examples/new_york_city.py b/examples/new_york_city.py index 8c0da7f..06c0a5f 100644 --- a/examples/new_york_city.py +++ b/examples/new_york_city.py @@ -12,6 +12,7 @@ bounds=(-74.26, 40.49, -73.70, 40.92), crs='EPSG:32618', features=['buildings', 'roads', 'water', 'fire', 'restaurant_grades', 'gtfs'], + weather=True, ao_samples=1, tour=tour, ) diff --git a/examples/nyc_lidar.py b/examples/nyc_lidar.py index fdfd0c8..32f8ab3 100644 --- a/examples/nyc_lidar.py +++ b/examples/nyc_lidar.py @@ -99,8 +99,16 @@ def main(): # Save all meshes (buildings + lidar) to zarr for next run terrain.rtx.save_meshes(zarr_path) + # Fetch weather for cloud + rain overlay (Shift+N) + weather = None + try: + from rtxpy import fetch_weather + weather = fetch_weather(dem_bounds, grid_size=15) + except Exception as e: + print(f"Skipping weather: {e}") + # Launch interactive viewer - terrain.rtx.explore(width=1600, height=1200) + terrain.rtx.explore(width=1600, height=1200, weather_data=weather) if __name__ == '__main__': diff --git a/examples/playground.py b/examples/playground.py index 6e5f93c..ab86fdb 100644 --- a/examples/playground.py +++ b/examples/playground.py @@ -236,6 +236,14 @@ def load_terrain(): except Exception as e: print(f"Skipping wind: {e}") + # --- Weather data (clouds + rain via Shift+N) ------------------------- + weather = None + try: + from rtxpy import fetch_weather + weather = fetch_weather(BOUNDS, grid_size=15) + except Exception as e: + print(f"Skipping weather: {e}") + # --- Hydro flow data (off by default, Shift+Y to toggle) --------------- hydro = None try: @@ -317,7 +325,7 @@ def load_terrain(): print(f"Skipping hydro: {e}") print("\nLaunching explore (press G to cycle layers, " - "Shift+W for wind, Shift+Y for hydro)...\n") + "Shift+W for wind, Shift+N for clouds, Shift+Y for hydro)...\n") ds.rtx.explore( z='elevation', scene_zarr=ZARR, @@ -326,6 +334,7 @@ def load_terrain(): height=768, render_scale=0.5, wind_data=wind, + weather_data=weather, hydro_data=hydro, fog_density=3.0, repl=True, diff --git a/examples/port_jefferson.py b/examples/port_jefferson.py index 89190c1..b8cdc12 100644 --- a/examples/port_jefferson.py +++ b/examples/port_jefferson.py @@ -22,6 +22,7 @@ source='usgs_1m', features=['buildings', 'roads', 'water'], tiles='satellite', + tile_zoom=18, hydro=True, wind=True, weather=True, diff --git a/examples/trinidad_tobago_coastal_resilience.py b/examples/trinidad_tobago_coastal_resilience.py index 71e57a1..d85e06f 100644 --- a/examples/trinidad_tobago_coastal_resilience.py +++ b/examples/trinidad_tobago_coastal_resilience.py @@ -381,6 +381,14 @@ def classify_roads(road_data, raw_elev): except Exception as e: print(f"Skipping wind: {e}") + # --- Weather (clouds + rain via Shift+N) ------------------------------- + weather = None + try: + from rtxpy import fetch_weather + weather = fetch_weather(BOUNDS, grid_size=15) + except Exception as e: + print(f"Skipping weather: {e}") + # --- Launch explorer --------------------------------------------------- print("\n" + "=" * 60) print("COASTAL RESILIENCE EXPLORER") @@ -388,6 +396,7 @@ def classify_roads(road_data, raw_elev): print(" G cycle layers (elevation → slope → surge_risk → flood maps)") print(" U toggle satellite overlay") print(" Shift+W toggle wind particles (storm simulation)") + print(" Shift+N toggle clouds + rain") print(" O / V set observer / toggle viewshed") print(" M minimap") print(" H full help overlay") @@ -414,6 +423,7 @@ def classify_roads(road_data, raw_elev): color_stretch='cbrt', subsample=1, wind_data=wind, + weather_data=weather, minimap_style='cyberpunk', minimap_layer='surge_risk', minimap_colors=RISK_COLORS, diff --git a/rtxpy/analysis/render.py b/rtxpy/analysis/render.py index c874a8a..ff31593 100644 --- a/rtxpy/analysis/render.py +++ b/rtxpy/analysis/render.py @@ -809,6 +809,282 @@ def _generate_reflection_rays(reflection_rays, primary_rays, primary_hits, ) +@cuda.jit(device=True, inline='always') +def _cloud_hash_i(n): + """Integer hash for procedural noise (works under Numba CUDA).""" + n = (n << 13) ^ n + return (n * (n * n * 15731 + 789221) + 1376312589) & 0x7fffffff + + +@cuda.jit(device=True, inline='always') +def _cloud_hash_3(a, b, c): + """Hash 3 integer coords to a float in [0, 1).""" + return _cloud_hash_i(a + _cloud_hash_i(b + _cloud_hash_i(c))) / 2147483647.0 + + +@cuda.jit(device=True, inline='always') +def _cloud_noise_3d(x, y, z): + """3D value noise via integer hashing + trilinear interpolation.""" + ix = int(math.floor(x)) + iy = int(math.floor(y)) + iz = int(math.floor(z)) + fx = x - ix + fy = y - iy + fz = z - iz + # Smoothstep + ux = fx * fx * (3.0 - 2.0 * fx) + uy = fy * fy * (3.0 - 2.0 * fy) + uz = fz * fz * (3.0 - 2.0 * fz) + # Hash 8 corners + n000 = _cloud_hash_3(ix, iy, iz) + n100 = _cloud_hash_3(ix + 1, iy, iz) + n010 = _cloud_hash_3(ix, iy + 1, iz) + n110 = _cloud_hash_3(ix + 1, iy + 1, iz) + n001 = _cloud_hash_3(ix, iy, iz + 1) + n101 = _cloud_hash_3(ix + 1, iy, iz + 1) + n011 = _cloud_hash_3(ix, iy + 1, iz + 1) + n111 = _cloud_hash_3(ix + 1, iy + 1, iz + 1) + # Trilinear + nx00 = n000 * (1.0 - ux) + n100 * ux + nx10 = n010 * (1.0 - ux) + n110 * ux + nx01 = n001 * (1.0 - ux) + n101 * ux + nx11 = n011 * (1.0 - ux) + n111 * ux + nxy0 = nx00 * (1.0 - uy) + nx10 * uy + nxy1 = nx01 * (1.0 - uy) + nx11 * uy + return nxy0 * (1.0 - uz) + nxy1 * uz + + +@cuda.jit(device=True, inline='always') +def _cloud_fbm_3d(x, y, z, octaves): + """Fractal Brownian motion: sum of noise at increasing frequencies.""" + value = 0.0 + amplitude = 0.5 + freq = 1.0 + for _ in range(octaves): + value += amplitude * _cloud_noise_3d(x * freq, y * freq, z * freq) + freq *= 2.0 + amplitude *= 0.5 + return value + + +@cuda.jit +def _volumetric_cloud_kernel( + output, primary_rays, primary_hits, + cloud_cover_grid, + sun_dir, + width, height, + cloud_base_z, cloud_top_z, + pixel_spacing_x, pixel_spacing_y, + grid_rows, grid_cols, + time_offset, + extinction, + terrain_rows, terrain_cols, +): + """Ray-march volumetric clouds and composite over shaded frame.""" + idx = cuda.grid(1) + if idx >= width * height: + return + + py = idx // width + px = idx % width + + # Read ray origin and direction + ox = primary_rays[idx, 0] + oy = primary_rays[idx, 1] + oz = primary_rays[idx, 2] + dx = primary_rays[idx, 4] + dy = primary_rays[idx, 5] + dz = primary_rays[idx, 6] + + # Hit distance (terrain/geometry); 0 or NaN means miss (sky) + hit_t = primary_hits[idx, 0] + # NaN-safe: treat NaN as miss (no terrain to occlude clouds) + if not (hit_t > 0.0): + hit_t = 0.0 + + # Ray-slab intersection: cloud_base_z to cloud_top_z (horizontal slab) + if abs(dz) < 1e-8: + return # Ray parallel to slab — skip + + t_base = (cloud_base_z - oz) / dz + t_top = (cloud_top_z - oz) / dz + entry_t = min(t_base, t_top) + exit_t = max(t_base, t_top) + + # Clamp to valid range + if entry_t < 0.0: + entry_t = 0.0 + # Clouds behind terrain are occluded + if hit_t > 0.0 and exit_t > hit_t: + exit_t = hit_t + if entry_t >= exit_t: + return + + slab_thickness = cloud_top_z - cloud_base_z + inv_slab = 1.0 / slab_thickness + + # Pre-loop: sample coverage + noise ONCE at ray midpoint. + # This eliminates all hash/noise from the inner loop. + mid_t = (entry_t + exit_t) * 0.5 + mid_x = ox + mid_t * dx + mid_y = oy + mid_t * dy + mid_gx = mid_x / pixel_spacing_x * grid_cols / terrain_cols + mid_gy = mid_y / pixel_spacing_y * grid_rows / terrain_rows + mid_ix = int(mid_gx) + mid_iy = int(mid_gy) + if mid_ix < 0 or mid_ix >= grid_cols or mid_iy < 0 or mid_iy >= grid_rows: + return + coverage = cloud_cover_grid[mid_iy, mid_ix] + if coverage < 0.05: + return + + # Single noise sample for cloud shape (2D at slab midpoint) + noise_scale = 8.0 * inv_slab + nx = mid_x * noise_scale + time_offset * 0.3 + ny = mid_y * noise_scale + time_offset * 0.1 + noise_val = _cloud_noise_3d(nx, ny, 0.0) + + # Pre-compute base shape: coverage * noise, then threshold + base_shape = coverage * noise_val + if base_shape < 0.05: + return + + # March through slab in 8 steps — inner loop is now trivially cheap + num_steps = 8 + step_size = (exit_t - entry_t) / num_steps + + # Sun direction for lighting + sun_x = sun_dir[0] + sun_y = sun_dir[1] + sun_z = sun_dir[2] + + # Henyey-Greenstein phase function (cos angle between view and sun) + cos_theta = -(dx * sun_x + dy * sun_y + dz * sun_z) + g = 0.6 + g2 = g * g + phase = 0.25 / 3.14159 * (1.0 - g2) / ((1.0 + g2 - 2.0 * g * cos_theta) ** 1.5) + g_back = -0.3 + g_back2 = g_back * g_back + phase_back = 0.25 / 3.14159 * (1.0 - g_back2) / ((1.0 + g_back2 - 2.0 * g_back * cos_theta) ** 1.5) + phase_total = 0.7 * phase + 0.3 * phase_back + + transmittance = 1.0 + cloud_r = 0.0 + cloud_g = 0.0 + cloud_b = 0.0 + + # Sun color: warm white + sun_cr = 1.0 + sun_cg = 0.95 + sun_cb = 0.85 + # Ambient sky light + amb_cr = 0.55 + amb_cg = 0.65 + amb_cb = 0.80 + + for step in range(num_steps): + if transmittance < 0.01: + break + + t = entry_t + (step + 0.5) * step_size + sz = oz + t * dz + + # Height fraction within slab [0,1] + h_frac = (sz - cloud_base_z) * inv_slab + if h_frac < 0.0 or h_frac > 1.0: + continue + + # Height gradient: cumulus profile + if h_frac < 0.3: + h_grad = h_frac / 0.3 + elif h_frac < 0.6: + h_grad = 1.0 + else: + h_grad = 1.0 - (h_frac - 0.6) * 2.5 + if h_grad <= 0.0: + continue + + # Density from pre-computed shape * height gradient + density = (base_shape * h_grad - 0.12) * 3.5 + if density <= 0.0: + continue + if density > 1.0: + density = 1.0 + + # Beer-Lambert + powder approximation + optical_step = density * step_size * extinction + beer = math.exp(-optical_step) + powder = 1.0 - math.exp(-2.0 * optical_step) + energy = 2.0 * beer * powder + + # Light energy: phase-function sun + ambient sky + light_r = energy * phase_total * sun_cr + energy * 0.35 * amb_cr + light_g = energy * phase_total * sun_cg + energy * 0.35 * amb_cg + light_b = energy * phase_total * sun_cb + energy * 0.35 * amb_cb + + # Accumulate + cloud_r += transmittance * light_r + cloud_g += transmittance * light_g + cloud_b += transmittance * light_b + transmittance *= beer + + # If no cloud contribution, skip + if transmittance > 0.999: + return + + # Composite: final = cloud_accumulated + transmittance * existing_pixel + existing_r = output[py, px, 0] + existing_g = output[py, px, 1] + existing_b = output[py, px, 2] + + output[py, px, 0] = cloud_r + transmittance * existing_r + output[py, px, 1] = cloud_g + transmittance * existing_g + output[py, px, 2] = cloud_b + transmittance * existing_b + + +_DUMMY_CLOUD_1x1 = None + + +def _apply_volumetric_clouds(d_output, d_primary_rays, d_primary_hits, + cloud_cover_grid, d_sun_dir, + width, height, + cloud_base_z, cloud_top_z, + pixel_spacing_x, pixel_spacing_y, + terrain_rows, terrain_cols, + cloud_time, extinction=0.8): + """Launch volumetric cloud ray march kernel.""" + global _DUMMY_CLOUD_1x1 + if cloud_cover_grid is None: + if _DUMMY_CLOUD_1x1 is None: + _DUMMY_CLOUD_1x1 = cupy.zeros((1, 1), dtype=np.float32) + cloud_cover_grid = _DUMMY_CLOUD_1x1 + grid_rows, grid_cols = cloud_cover_grid.shape + if grid_rows <= 1: + return # No real cloud data + # Normalize extinction by slab thickness so visual density is + # scale-independent. Target ~3 optical depths at full density + # across the entire slab → translucent-to-opaque clouds. + slab_thickness = cloud_top_z - cloud_base_z + if slab_thickness > 0: + extinction = 3.0 / slab_thickness + + num_pixels = width * height + threadsperblock = 256 + blockspergrid = (num_pixels + threadsperblock - 1) // threadsperblock + _volumetric_cloud_kernel[blockspergrid, threadsperblock]( + d_output, d_primary_rays, d_primary_hits, + cloud_cover_grid, + d_sun_dir, + np.int32(width), np.int32(height), + np.float32(cloud_base_z), np.float32(cloud_top_z), + np.float32(pixel_spacing_x), np.float32(pixel_spacing_y), + np.int32(grid_rows), np.int32(grid_cols), + np.float32(cloud_time), + np.float32(extinction), + np.int32(terrain_rows), np.int32(terrain_cols), + ) + + @cuda.jit def _shade_terrain_kernel( output, albedo_out, primary_rays, primary_hits, shadow_hits, @@ -827,7 +1103,8 @@ def _shade_terrain_kernel( instance_ids, geometry_colors, primitive_ids, point_colors, point_color_offsets, ao_factor, gi_color, gi_intensity, - reflection_hits, reflection_rays + reflection_hits, reflection_rays, + cloud_fog_map, cloud_fog_density ): """GPU kernel for terrain shading with lighting, shadows, fog, colormapping, and viewshed.""" idx = cuda.grid(1) @@ -1051,6 +1328,37 @@ def _shade_terrain_kernel( if shadow_t > 0: shadow_factor = 0.5 + # Cloud shadow — patchy darkening where cloud_cover is high + cfm_h = cloud_fog_map.shape[0] + cfm_w = cloud_fog_map.shape[1] + cloud_shadow = 0.0 + if cfm_h > 1: + cy = elev_y + if cy < 0: + cy = 0 + elif cy >= cfm_h: + cy = cfm_h - 1 + cx = elev_x + if cx < 0: + cx = 0 + elif cx >= cfm_w: + cx = cfm_w - 1 + cc = cloud_fog_map[cy, cx] + if cc > 0.01: + # Procedural noise for patchy cloud shapes + wx = hit_x * cloud_fog_density + wy = hit_y * cloud_fog_density + n = (math.sin(wx * 1.2 + wy * 0.7) * 0.45 + + math.sin(wx * 2.7 - wy * 1.8 + 3.1) * 0.30 + + math.sin(wx * 0.5 + wy * 3.4 - 1.7) * 0.25) + # n in ~[-1,1] → [0,1], squared for sharper edges + cloud_mask = n * 0.5 + 0.5 + cloud_mask = cloud_mask * cloud_mask + cloud_shadow = cc * cloud_mask * 0.7 + if cloud_shadow > 0.7: + cloud_shadow = 0.7 + shadow_factor *= (1.0 - cloud_shadow) + # Final lighting diffuse = cos_theta * shadow_factor lighting = ambient + (1.0 - ambient) * diffuse @@ -1235,6 +1543,16 @@ def _shade_terrain_kernel( color_g = color_g * (1 - fog_amount) + fog_color_g * fog_amount color_b = color_b * (1 - fog_amount) + fog_color_b * fog_amount + # Cloud desaturation — overcast areas lose color contrast + if cloud_shadow > 0.15: + gray = color_r * 0.3 + color_g * 0.59 + color_b * 0.11 + desat = (cloud_shadow - 0.15) * 0.6 + if desat > 0.35: + desat = 0.35 + color_r = color_r * (1.0 - desat) + gray * desat + color_g = color_g * (1.0 - desat) + gray * desat + color_b = color_b * (1.0 - desat) + gray * desat + output[py, px, 0] = color_r output[py, px, 1] = color_g output[py, px, 2] = color_b @@ -1580,6 +1898,7 @@ def _shade_terrain( ao_factor=None, gi_color=None, gi_intensity=2.0, reflection_hits=None, reflection_rays=None, albedo_out=None, + cloud_fog_map=None, cloud_fog_density=0.0, ): """Apply terrain shading with all effects.""" threadsperblock = 256 @@ -1674,6 +1993,13 @@ def _shade_terrain( _DUMMY_ALBEDO = cupy.zeros((1, 1, 3), dtype=np.float32) albedo_out = _DUMMY_ALBEDO + # Handle cloud fog map - dummy (1,1) when not provided + # Kernel checks shape[0] > 1 to decide whether to apply + if cloud_fog_map is None: + if _DUMMY_1x1 is None: + _DUMMY_1x1 = cupy.zeros((1, 1), dtype=np.float32) + cloud_fog_map = _DUMMY_1x1 + _shade_terrain_kernel[blockspergrid, threadsperblock]( output, albedo_out, primary_rays, primary_hits, shadow_hits, elevation_data, color_lut, num_rays, width, height, @@ -1691,7 +2017,8 @@ def _shade_terrain( instance_ids, geometry_colors, primitive_ids, point_colors, point_color_offsets, ao_factor, gi_color, np.float32(gi_intensity), - reflection_hits, reflection_rays + reflection_hits, reflection_rays, + cloud_fog_map, np.float32(cloud_fog_density) ) @@ -1840,6 +2167,12 @@ def render( edl: bool = True, edl_strength: float = 0.7, edl_radius: float = 2.0, + cloud_fog_map=None, + cloud_fog_density: float = 0.0, + volumetric_clouds: bool = False, + cloud_base_z: float = 0.0, + cloud_top_z: float = 0.0, + cloud_time: float = 0.0, _return_gpu: bool = False, ) -> np.ndarray: """Render terrain with a perspective camera for movie-quality visualization. @@ -2177,6 +2510,14 @@ def render( ov_max = float(cupy.nanmax(d_overlay)) ov_range = ov_max - ov_min + # Prepare cloud fog map for GPU + d_cloud_fog_map = None + if cloud_fog_map is not None: + if not isinstance(cloud_fog_map, cupy.ndarray): + d_cloud_fog_map = cupy.asarray(cloud_fog_map, dtype=cupy.float32) + else: + d_cloud_fog_map = cloud_fog_map if cloud_fog_map.dtype == cupy.float32 else cloud_fog_map.astype(cupy.float32) + # Build per-point color buffers for sphere geometries d_point_colors = None d_point_color_offsets = None @@ -2209,8 +2550,24 @@ def render( gi_intensity=gi_intensity, reflection_hits=d_reflection_hits, reflection_rays=d_reflection_rays, albedo_out=bufs.albedo, + cloud_fog_map=None if volumetric_clouds else d_cloud_fog_map, + cloud_fog_density=cloud_fog_density, ) + # Volumetric clouds (ray-marched, composited over shaded frame) + if volumetric_clouds and cloud_top_z > cloud_base_z: + d_primary_rays_flat = bufs.primary_rays + d_primary_hits_flat = bufs.primary_hits + _apply_volumetric_clouds( + d_output, d_primary_rays_flat, d_primary_hits_flat, + d_cloud_fog_map, d_sun_dir, + width, height, + cloud_base_z, cloud_top_z, + pixel_spacing_x, pixel_spacing_y, + H, W, + cloud_time, + ) + # AI denoiser (after shading, before bloom/tone mapping) if denoise: from ..rtx import denoise as _denoise diff --git a/rtxpy/engine.py b/rtxpy/engine.py index d17bc58..fd7904b 100644 --- a/rtxpy/engine.py +++ b/rtxpy/engine.py @@ -27,6 +27,7 @@ from .viewer.terrain import TerrainState from .viewer.observers import ObserverManager, Observer, OBSERVER_COLORS from .viewer.wind import WindState +from .viewer.cloud import CloudState from .viewer.hydro import HydroState from .viewer.hud import HUDState from .viewer.keybindings import MOVEMENT_KEYS, SHIFT_BINDINGS, KEY_BINDINGS, SPECIAL_BINDINGS @@ -321,6 +322,94 @@ def _hydro_splat_kernel( cuda.atomic.add(output, (py, px, 2), contrib * color_b) + @cuda.jit + def _rain_splat_kernel( + pts, # (N, 2) float32 — (row, col) per rain particle + z_frac, # (N,) float32 — altitude fraction (0=ground, 1=cloud) + alphas, # (N,) float32 — pre-computed alpha + streak_lens, # (N,) int32 — vertical streak length in pixels + terrain, # (tH, tW) float32 — terrain elevation + output, # (sh, sw, 3) float32 — frame buffer (atomic add) + # Camera basis + cam_x, cam_y, cam_z, + fwd_x, fwd_y, fwd_z, + rgt_x, rgt_y, rgt_z, + up_x, up_y, up_z, + # Projection + fov_scale, aspect_ratio, + # World params + psx, psy, ve, subsample_f, cloud_z, min_depth, + # Color + color_r, color_g, color_b, + ): + idx = cuda.grid(1) + if idx >= pts.shape[0]: + return + + a = alphas[idx] + if a < 0.002: + return + + row = pts[idx, 0] + col = pts[idx, 1] + + # Terrain Z lookup + tH = terrain.shape[0] + tW = terrain.shape[1] + sr = int(row / subsample_f) + sc = int(col / subsample_f) + if sr < 0: + sr = 0 + elif sr >= tH: + sr = tH - 1 + if sc < 0: + sc = 0 + elif sc >= tW: + sc = tW - 1 + z_raw = terrain[sr, sc] + if z_raw != z_raw: + z_raw = 0.0 + terrain_z = z_raw * ve + rain_z = terrain_z + z_frac[idx] * (cloud_z - terrain_z) + + wx = col * psx + wy = row * psy + + dx = wx - cam_x + dy = wy - cam_y + dz = rain_z - cam_z + + depth = dx * fwd_x + dy * fwd_y + dz * fwd_z + if depth <= min_depth: + return + + inv_depth = 1.0 / (depth + 1e-10) + u_cam = dx * rgt_x + dy * rgt_y + dz * rgt_z + v_cam = dx * up_x + dy * up_y + dz * up_z + u_ndc = u_cam * inv_depth / (fov_scale * aspect_ratio) + v_ndc = v_cam * inv_depth / fov_scale + + sh = output.shape[0] + sw = output.shape[1] + sx = int((u_ndc + 1.0) * 0.5 * sw) + sy = int((1.0 - v_ndc) * 0.5 * sh) + + if sx < 0 or sx >= sw or sy < 0 or sy >= sh: + return + + # Vertical streak + sl = streak_lens[idx] + for dy_off in range(sl): + py = sy + dy_off + if py < 0 or py >= sh: + continue + t = float(dy_off) / float(sl) if sl > 0 else 0.0 + streak_a = a * (1.0 - t * 0.6) + cuda.atomic.add(output, (py, sx, 0), streak_a * color_r) + cuda.atomic.add(output, (py, sx, 1), streak_a * color_g) + cuda.atomic.add(output, (py, sx, 2), streak_a * color_b) + + def _glfw_to_key(glfw_key, mods): """Translate a GLFW key code + modifiers to the string format used by _handle_key_press / _handle_key_release. @@ -1459,6 +1548,9 @@ def __init__(self, raster, width: int = 800, height: int = 600, # Wind particle state self.wind = WindState() + # Cloud particle state + self.clouds = CloudState() + # Hydro flow particle state self.hydro = HydroState() @@ -2407,6 +2499,130 @@ def _wind_done_event(self): def _wind_done_event(self, value): self.wind.wind_done_event = value + # ------------------------------------------------------------------ + # Delegation properties — CloudState + # ------------------------------------------------------------------ + + @property + def _clouds_enabled(self): + return self.clouds.clouds_enabled + + @_clouds_enabled.setter + def _clouds_enabled(self, value): + self.clouds.clouds_enabled = value + + @property + def _cloud_cover_grid(self): + return self.clouds.cloud_cover_grid + + @_cloud_cover_grid.setter + def _cloud_cover_grid(self, value): + self.clouds.cloud_cover_grid = value + + @property + def _cloud_particles(self): + return self.clouds.cloud_particles + + @_cloud_particles.setter + def _cloud_particles(self, value): + self.clouds.cloud_particles = value + + @property + def _cloud_sizes(self): + return self.clouds.cloud_sizes + + @_cloud_sizes.setter + def _cloud_sizes(self, value): + self.clouds.cloud_sizes = value + + @property + def _cloud_alphas(self): + return self.clouds.cloud_alphas + + @_cloud_alphas.setter + def _cloud_alphas(self, value): + self.clouds.cloud_alphas = value + + @property + def _cloud_ages(self): + return self.clouds.cloud_ages + + @_cloud_ages.setter + def _cloud_ages(self, value): + self.clouds.cloud_ages = value + + @property + def _cloud_lifetimes(self): + return self.clouds.cloud_lifetimes + + @_cloud_lifetimes.setter + def _cloud_lifetimes(self, value): + self.clouds.cloud_lifetimes = value + + @property + def _cloud_n_particles(self): + return self.clouds.cloud_n_particles + + @_cloud_n_particles.setter + def _cloud_n_particles(self, value): + self.clouds.cloud_n_particles = value + + @property + def _cloud_max_age(self): + return self.clouds.cloud_max_age + + @_cloud_max_age.setter + def _cloud_max_age(self, value): + self.clouds.cloud_max_age = value + + @property + def _cloud_altitude(self): + return self.clouds.cloud_altitude + + @_cloud_altitude.setter + def _cloud_altitude(self, value): + self.clouds.cloud_altitude = value + + @property + def _cloud_min_depth(self): + return self.clouds.cloud_min_depth + + @_cloud_min_depth.setter + def _cloud_min_depth(self, value): + self.clouds.cloud_min_depth = value + + @property + def _cloud_terrain_np(self): + return self.clouds.cloud_terrain_np + + @_cloud_terrain_np.setter + def _cloud_terrain_np(self, value): + self.clouds.cloud_terrain_np = value + + @property + def _volumetric_clouds_enabled(self): + return self.clouds.volumetric_clouds_enabled + + @_volumetric_clouds_enabled.setter + def _volumetric_clouds_enabled(self, value): + self.clouds.volumetric_clouds_enabled = value + + @property + def _cloud_thickness(self): + return self.clouds.cloud_thickness + + @_cloud_thickness.setter + def _cloud_thickness(self, value): + self.clouds.cloud_thickness = value + + @property + def _cloud_time(self): + return self.clouds.cloud_time + + @_cloud_time.setter + def _cloud_time(self, value): + self.clouds.cloud_time = value + # ------------------------------------------------------------------ # Delegation properties — HydroState # ------------------------------------------------------------------ @@ -3430,6 +3646,7 @@ def _rebuild_at_resolution(self, factor): self._hydro_terrain_np = None self._d_base_frame = None # invalidate GPU wind/hydro buffers self._d_wind_scratch = None + self._d_cloud_fog_map = None # invalidate cached cloud fog map H, W = sub.shape self.terrain_shape = (H, W) @@ -4316,6 +4533,378 @@ def _init_weather(self, weather_data): if added: print(f" Weather: {len(added)} layers added ({', '.join(added)})") + # Auto-initialize cloud + rain particles from weather data + self._init_clouds(weather_data) + + # ------------------------------------------------------------------ + # Cloud + rain particle system + # ------------------------------------------------------------------ + + def _init_clouds(self, weather_data): + """Initialize cloud and rain particles from weather data. + + Spawns cloud puffs proportional to cloud_cover and rain streaks + proportional to precipitation. Both drift with the wind field + if available. + """ + if weather_data is None: + return + + cloud_cover = weather_data.get('cloud_cover') + precipitation = weather_data.get('precipitation') + if cloud_cover is None and precipitation is None: + return + + from .tiles import _build_latlon_grids + from scipy.interpolate import RegularGridInterpolator + + raster = self._base_raster + H, W = raster.shape + w_lats = weather_data['lats'] + w_lons = weather_data['lons'] + + lats_grid, lons_grid = _build_latlon_grids(raster) + points = np.stack([lats_grid.ravel(), lons_grid.ravel()], axis=-1) + + # Interpolate cloud_cover (0-100%) to terrain grid + if cloud_cover is not None: + interp = RegularGridInterpolator( + (w_lats, w_lons), cloud_cover.astype(np.float32), + method='linear', bounds_error=False, fill_value=0.0, + ) + cc = interp(points).reshape(H, W).astype(np.float32) + self._cloud_cover_grid = np.clip(cc / 100.0, 0, 1) + else: + self._cloud_cover_grid = np.zeros((H, W), dtype=np.float32) + + # Interpolate precipitation (mm) to terrain grid + if precipitation is not None: + interp_p = RegularGridInterpolator( + (w_lats, w_lons), precipitation.astype(np.float32), + method='linear', bounds_error=False, fill_value=0.0, + ) + self._rain_grid = interp_p(points).reshape(H, W).astype(np.float32) + else: + self._rain_grid = np.zeros((H, W), dtype=np.float32) + + # Cloud altitude: above terrain max + terrain_data = raster.data + if hasattr(terrain_data, 'get'): + terrain_np = terrain_data.get() + else: + terrain_np = np.asarray(terrain_data) + psx = self._base_pixel_spacing_x + psy = self._base_pixel_spacing_y + z_max = float(np.nanmax(terrain_np)) + z_range = z_max - float(np.nanmin(terrain_np)) + diag = np.sqrt((W * psx) ** 2 + (H * psy) ** 2) + self._cloud_altitude = z_max + max(z_range * 0.5, diag * 0.04) + self._cloud_min_depth = diag * 0.01 + self._cloud_thickness = max(z_range * 0.3, diag * 0.02) + self._volumetric_clouds_enabled = True + + # Spawn cloud particles weighted by cloud_cover density + N = self._cloud_n_particles + cc_flat = self._cloud_cover_grid.ravel() + cc_sum = cc_flat.sum() + if cc_sum > 0: + prob = cc_flat / cc_sum + else: + prob = np.ones_like(cc_flat) / cc_flat.size + + chosen = np.random.choice(cc_flat.size, size=N, p=prob) + rows = (chosen // W).astype(np.float32) + cols = (chosen % W).astype(np.float32) + # Jitter within cell + rows += np.random.uniform(-0.5, 0.5, N).astype(np.float32) + cols += np.random.uniform(-0.5, 0.5, N).astype(np.float32) + self._cloud_particles = np.column_stack([rows, cols]).astype(np.float32) + + # Per-particle size (world-space radius) + base_size = max(psx, psy) * 8 + self._cloud_sizes = ( + np.random.uniform(0.6, 1.8, N).astype(np.float32) * base_size + ) + + # Per-particle alpha from local cloud_cover + r_idx = np.clip(rows.astype(int), 0, H - 1) + c_idx = np.clip(cols.astype(int), 0, W - 1) + local_cc = self._cloud_cover_grid[r_idx, c_idx] + self._cloud_alphas = (local_cc * 0.35).astype(np.float32) + + # Ages / lifetimes + self._cloud_max_age = 300 + self._cloud_lifetimes = np.random.randint( + self._cloud_max_age // 2, self._cloud_max_age, N) + self._cloud_ages = np.random.randint(0, self._cloud_max_age, N) + + # Rain particles — fewer, driven by precipitation + rain_N = 8000 + rain_flat = self._rain_grid.ravel() + rain_sum = rain_flat.sum() + if rain_sum > 0: + rain_prob = rain_flat / rain_sum + chosen_r = np.random.choice(rain_flat.size, size=rain_N, p=rain_prob) + else: + chosen_r = np.random.randint(0, rain_flat.size, rain_N) + r_rows = (chosen_r // W).astype(np.float32) + r_cols = (chosen_r % W).astype(np.float32) + r_rows += np.random.uniform(-0.5, 0.5, rain_N).astype(np.float32) + r_cols += np.random.uniform(-0.5, 0.5, rain_N).astype(np.float32) + self._rain_particles = np.column_stack([r_rows, r_cols]).astype(np.float32) + # Rain z: random altitude between terrain and cloud layer + self._rain_z_frac = np.random.uniform(0.0, 1.0, rain_N).astype(np.float32) + self._rain_ages = np.random.randint(0, 40, rain_N) + self._rain_lifetimes = np.random.randint(20, 40, rain_N) + self._rain_max_age = 40 + + mean_cc = float(np.mean(self._cloud_cover_grid) * 100) + mean_precip = float(np.mean(self._rain_grid)) + print(f" Clouds: {N} particles (mean cover {mean_cc:.0f}%)" + f" Rain: {rain_N} particles (mean precip {mean_precip:.1f}mm)") + + def _toggle_clouds(self): + """Toggle cloud fog + rain on/off.""" + if self._cloud_cover_grid is None: + print("No weather data loaded. Pass weather_data to explore().") + return + self._clouds_enabled = not self._clouds_enabled + if self._clouds_enabled and self._cloud_cover_grid is not None: + cc = self._cloud_cover_grid + print(f"Clouds: ON (cover min={cc.min():.2f} max={cc.max():.2f} mean={cc.mean():.2f})") + else: + print(f"Clouds: OFF") + self._render_needed = True + + def _action_toggle_clouds(self): + self._toggle_clouds() + + def _update_rain_particles(self): + """Animate rain streaks: fall downward and respawn.""" + if not hasattr(self, '_rain_particles') or self._rain_particles is None: + return + + pts = self._rain_particles + H, W = self._cloud_cover_grid.shape + + # Rain falls: decrease z_frac (toward terrain) + self._rain_z_frac -= 0.06 + self._rain_ages += 1 + + # Slight wind drift + if self._wind_u_px is not None: + r0 = np.clip(pts[:, 0].astype(int), 0, H - 1) + c0 = np.clip(pts[:, 1].astype(int), 0, W - 1) + s = getattr(self, '_dt_scale', 1.0) + pts[:, 1] += self._wind_u_px[r0, c0] * 0.15 * s + pts[:, 0] += self._wind_v_px[r0, c0] * 0.15 * s + + # Respawn when hitting ground or aged out + respawn = ( + (self._rain_z_frac <= 0) | + (self._rain_ages >= self._rain_lifetimes) | + (pts[:, 0] < 0) | (pts[:, 0] >= H) | + (pts[:, 1] < 0) | (pts[:, 1] >= W) + ) + n_respawn = int(respawn.sum()) + if n_respawn > 0: + rain_flat = self._rain_grid.ravel() + rain_sum = rain_flat.sum() + if rain_sum > 0: + prob = rain_flat / rain_sum + else: + prob = np.ones_like(rain_flat) / rain_flat.size + chosen = np.random.choice(rain_flat.size, size=n_respawn, p=prob) + pts[respawn, 0] = (chosen // W).astype(np.float32) + np.random.uniform(-0.5, 0.5, n_respawn) + pts[respawn, 1] = (chosen % W).astype(np.float32) + np.random.uniform(-0.5, 0.5, n_respawn) + self._rain_z_frac[respawn] = np.random.uniform(0.7, 1.0, n_respawn).astype(np.float32) + self._rain_ages[respawn] = 0 + self._rain_lifetimes[respawn] = np.random.randint(20, 40, n_respawn) + + def _draw_rain_on_frame(self, img, forward, right, cam_up, + fov_scale, aspect, min_depth): + """Render rain streaks on the frame.""" + sh, sw = img.shape[:2] + psx = self._base_pixel_spacing_x + psy = self._base_pixel_spacing_y + ve = self.vertical_exaggeration + cam_pos = self.position + + # Cached terrain for z lookup + if self._cloud_terrain_np is None: + terrain_data = self.raster.data + if hasattr(terrain_data, 'get'): + self._cloud_terrain_np = terrain_data.get() + else: + self._cloud_terrain_np = np.asarray(terrain_data) + terrain_np = self._cloud_terrain_np + tH, tW = terrain_np.shape + f = self.subsample_factor + cloud_z = self._cloud_altitude * ve + + pts = self._rain_particles + N = pts.shape[0] + + # Rain z: interpolate between terrain and cloud altitude + sr = np.clip((pts[:, 0] / f).astype(int), 0, tH - 1) + sc = np.clip((pts[:, 1] / f).astype(int), 0, tW - 1) + terrain_z = np.nan_to_num(terrain_np[sr, sc], nan=0.0) * ve + rain_z = terrain_z + self._rain_z_frac * (cloud_z - terrain_z) + + wx = pts[:, 1] * psx + wy = pts[:, 0] * psy + + dx = wx - cam_pos[0] + dy = wy - cam_pos[1] + dz = rain_z - cam_pos[2] + + depth = dx * forward[0] + dy * forward[1] + dz * forward[2] + valid = depth > min_depth + + inv_depth = np.where(valid, 1.0 / (depth + 1e-10), 0.0) + u_cam = dx * right[0] + dy * right[1] + dz * right[2] + v_cam = dx * cam_up[0] + dy * cam_up[1] + dz * cam_up[2] + u_ndc = u_cam * inv_depth / (fov_scale * aspect) + v_ndc = v_cam * inv_depth / fov_scale + + sx = ((u_ndc + 1.0) * 0.5 * sw).astype(np.int32) + sy = ((1.0 - v_ndc) * 0.5 * sh).astype(np.int32) + + # Rain streak length (vertical on screen, 3-8 pixels) + streak_len = np.clip((5.0 * inv_depth * sh * fov_scale * 0.003), 2, 8).astype(np.int32) + + # Alpha based on precipitation intensity + rain_r = np.clip((pts[:, 0] / f).astype(int), 0, tH - 1) + rain_c = np.clip((pts[:, 1] / f).astype(int), 0, tW - 1) + local_precip = self._rain_grid[ + np.clip(pts[:, 0].astype(int), 0, self._rain_grid.shape[0] - 1), + np.clip(pts[:, 1].astype(int), 0, self._rain_grid.shape[1] - 1), + ] + base_alpha = np.clip(local_precip / 5.0, 0.05, 0.5).astype(np.float32) + + # Fade + ages = self._rain_ages.astype(np.float32) + lifetimes = self._rain_lifetimes.astype(np.float32) + fade = np.clip(ages / 3.0, 0, 1) * np.clip((lifetimes - ages) / 5.0, 0, 1) + alpha = base_alpha * fade * 0.15 + + on_screen = ( + valid & + (sx >= 0) & (sx < sw) & + (sy >= 0) & (sy < sh + 8) & + (alpha > 0.002) + ) + if not on_screen.any(): + return + + sx_v = sx[on_screen] + sy_v = sy[on_screen] + al_v = alpha[on_screen] + sl_v = streak_len[on_screen] + + # Draw vertical streaks (light blue-gray) + color = np.array([0.7, 0.75, 0.85], dtype=np.float32) + max_sl = int(sl_v.max()) if sl_v.size > 0 else 3 + for dy_off in range(max_sl): + py = sy_v + dy_off + streak_ok = (dy_off < sl_v) & (py >= 0) & (py < sh) + if not streak_ok.any(): + continue + t = dy_off / (sl_v[streak_ok].astype(np.float32)) + streak_alpha = al_v[streak_ok] * (1.0 - t * 0.6) + for c in range(3): + np.add.at(img[:, :, c], (py[streak_ok], sx_v[streak_ok]), + streak_alpha * color[c]) + + np.clip(img, 0, 1, out=img) + + def _splat_rain_gpu(self, d_frame): + """Project and splat rain particles as vertical streaks on GPU.""" + if not hasattr(self, '_rain_particles') or self._rain_particles is None: + return + rain_grid = getattr(self, '_rain_grid', None) + if rain_grid is None or rain_grid.sum() < 0.01: + return + + from .analysis.render import _compute_camera_basis + + sh, sw = d_frame.shape[:2] + psx = float(self._base_pixel_spacing_x) + psy = float(self._base_pixel_spacing_y) + ve = float(self.vertical_exaggeration) + cloud_z = float(self._cloud_altitude * ve) + min_depth = float(self._cloud_min_depth) + + cam_pos = self.position + look_at = self._get_look_at() + forward, right, cam_up = _compute_camera_basis( + tuple(cam_pos), tuple(look_at), (0, 0, 1), + ) + fov_scale = float(math.tan(math.radians(self.fov) / 2.0)) + aspect = float(sw / sh) + + rain_pts = self._rain_particles + rain_N = rain_pts.shape[0] + f = float(self.subsample_factor) + + # GPU terrain for z lookup + terrain_data = self.raster.data + if not isinstance(terrain_data, cp.ndarray): + terrain_data = cp.asarray(terrain_data) + + # Pre-compute rain alpha + streak length on CPU + local_precip = self._rain_grid[ + np.clip(rain_pts[:, 0].astype(int), 0, self._rain_grid.shape[0] - 1), + np.clip(rain_pts[:, 1].astype(int), 0, self._rain_grid.shape[1] - 1), + ] + base_alpha = np.clip(local_precip / 5.0, 0.05, 0.5).astype(np.float32) + rain_ages = self._rain_ages.astype(np.float32) + rain_lifetimes = self._rain_lifetimes.astype(np.float32) + fade = (np.clip(rain_ages / 3.0, 0, 1) + * np.clip((rain_lifetimes - rain_ages) / 5.0, 0, 1)) + rain_alpha = (base_alpha * fade * 0.15).astype(np.float32) + + # Streak length: compute from depth + r_wx = rain_pts[:, 1] * psx + r_wy = rain_pts[:, 0] * psy + r_depth = ((r_wx - cam_pos[0]) * forward[0] + + (r_wy - cam_pos[1]) * forward[1] + + (cloud_z * 0.5 - cam_pos[2]) * forward[2]) + r_inv = np.where(r_depth > min_depth, 1.0 / (r_depth + 1e-10), 0.0) + streak = np.clip((5.0 * r_inv * sh * fov_scale * 0.003), 2, 8).astype(np.int32) + + # Upload + _dr = getattr(self, '_d_rain_pts', None) + if _dr is None or _dr.shape[0] != rain_N: + self._d_rain_pts = cp.empty((rain_N, 2), dtype=cp.float32) + self._d_rain_zfrac = cp.empty(rain_N, dtype=cp.float32) + self._d_rain_alpha = cp.empty(rain_N, dtype=cp.float32) + self._d_rain_streak = cp.empty(rain_N, dtype=cp.int32) + self._d_rain_pts.set(rain_pts) + self._d_rain_zfrac.set(self._rain_z_frac) + self._d_rain_alpha.set(rain_alpha) + self._d_rain_streak.set(streak) + + tpb = 256 + bpg_r = (rain_N + tpb - 1) // tpb + _rain_splat_kernel[bpg_r, tpb]( + self._d_rain_pts, + self._d_rain_zfrac, + self._d_rain_alpha, + self._d_rain_streak, + terrain_data, + d_frame, + float(cam_pos[0]), float(cam_pos[1]), float(cam_pos[2]), + float(forward[0]), float(forward[1]), float(forward[2]), + float(right[0]), float(right[1]), float(right[2]), + float(cam_up[0]), float(cam_up[1]), float(cam_up[2]), + fov_scale, aspect, + psx, psy, ve, f, cloud_z, min_depth, + 0.7, 0.75, 0.85, + ) + + cp.clip(d_frame, 0, 1, out=d_frame) + def _init_wind(self, wind_data): """Interpolate wind U/V from lat/lon grid onto the terrain pixel grid. @@ -6405,8 +6994,10 @@ def _tick(self): if self._render_needed: self._update_frame() self._render_needed = False - elif (self._wind_enabled or self._hydro_enabled) and self._d_base_frame is not None: - # Wind/hydro is on but camera didn't move — skip the expensive ray + elif ((self._wind_enabled or self._hydro_enabled + or (self._clouds_enabled and self._rain_particles is not None)) + and self._d_base_frame is not None): + # Particles active but camera didn't move — skip the expensive ray # trace and just re-advect particles + GPU splat on fresh copy. cp.copyto(self._d_wind_scratch, self._d_base_frame) if self._wind_enabled and self._wind_particles is not None: @@ -6415,6 +7006,9 @@ def _tick(self): if self._hydro_enabled and self._hydro_particles is not None: self._update_hydro_particles() self._splat_hydro_gpu(self._d_wind_scratch) + if self._clouds_enabled and self._rain_particles is not None: + self._update_rain_particles() + self._splat_rain_gpu(self._d_wind_scratch) self._d_wind_scratch.get(out=self._pinned_frame) self._composite_overlays() @@ -7386,6 +7980,22 @@ def _save_screenshot(self): # Common render kwargs # fog_density is scene-relative; convert to absolute for the kernel _fog = self.fog_density / self._scene_diagonal if self.fog_density > 0 else 0.0 + + # Cloud fog map for screenshot — reuse cached GPU array + _cloud_fog_map = None + _cloud_fog_density = 0.0 + if self._clouds_enabled and self._cloud_cover_grid is not None: + _d = getattr(self, '_d_cloud_fog_map', None) + f = self.subsample_factor + if _d is None or getattr(self, '_cloud_fog_subsample', 0) != f: + src = self._cloud_cover_grid + if f > 1: + src = src[::f, ::f] + self._d_cloud_fog_map = cp.asarray(src, dtype=cp.float32) + self._cloud_fog_subsample = f + _cloud_fog_map = self._d_cloud_fog_map + _cloud_fog_density = 12.0 / self._scene_diagonal + render_kwargs = dict( camera_position=tuple(self.position), look_at=tuple(self._get_look_at()), @@ -7413,6 +8023,12 @@ def _save_screenshot(self): overlay_as_water=self._overlay_as_water, overlay_color_lut=self._active_overlay_color_lut, geometry_colors=geometry_colors, + cloud_fog_map=_cloud_fog_map, + cloud_fog_density=_cloud_fog_density, + volumetric_clouds=self._volumetric_clouds_enabled and self._clouds_enabled, + cloud_base_z=self._cloud_altitude, + cloud_top_z=self._cloud_altitude + self._cloud_thickness, + cloud_time=self._cloud_time, ) # Accumulated multi-frame screenshot when AO or DOF is enabled @@ -7537,6 +8153,24 @@ def _render_frame(self): # fog_density is scene-relative; convert to absolute for the kernel _fog = self.fog_density / self._scene_diagonal if self.fog_density > 0 else 0.0 + + # Cloud fog map: pass cloud_cover_grid for cloud shadow modulation + _cloud_fog_map = None + _cloud_fog_density = 0.0 + if self._clouds_enabled and self._cloud_cover_grid is not None: + # Cache GPU array to avoid per-frame upload + _d = getattr(self, '_d_cloud_fog_map', None) + f = self.subsample_factor + if _d is None or getattr(self, '_cloud_fog_subsample', 0) != f: + src = self._cloud_cover_grid + if f > 1: + src = src[::f, ::f] + self._d_cloud_fog_map = cp.asarray(src, dtype=cp.float32) + self._cloud_fog_subsample = f + _cloud_fog_map = self._d_cloud_fog_map + # Spatial frequency: ~12 cloud cells across scene + _cloud_fog_density = 12.0 / self._scene_diagonal + d_output = render( self.raster, camera_position=tuple(self.position), @@ -7579,6 +8213,12 @@ def _render_frame(self): edge_strength=0.2, edge_color=(0.15, 0.13, 0.10), edl=self.edl_enabled, + cloud_fog_map=_cloud_fog_map, + cloud_fog_density=_cloud_fog_density, + volumetric_clouds=self._volumetric_clouds_enabled and self._clouds_enabled, + cloud_base_z=self._cloud_altitude, + cloud_top_z=self._cloud_altitude + self._cloud_thickness, + cloud_time=self._cloud_time, bloom=not defer_post, tone_map=not defer_post, _return_gpu=True, @@ -7697,8 +8337,10 @@ def _update_frame(self): self._pinned_mem, dtype=np.float32, count=d_display.size ).reshape(d_display.shape) - # Save clean post-processed frame for idle wind/hydro replay - if self._wind_enabled or self._hydro_enabled: + # Save clean post-processed frame for idle wind/hydro/rain replay + _any_particles = (self._wind_enabled or self._hydro_enabled + or (self._clouds_enabled and self._rain_particles is not None)) + if _any_particles: if self._d_base_frame is None or self._d_base_frame.shape != d_display.shape: self._d_base_frame = cp.empty_like(d_display) self._d_wind_scratch = cp.empty_like(d_display) @@ -7714,8 +8356,17 @@ def _update_frame(self): self._update_hydro_particles() self._splat_hydro_gpu(d_display) + # GPU rain: advect on CPU, splat on GPU (cloud fog is baked into ray trace) + if self._clouds_enabled and self._rain_particles is not None: + self._update_rain_particles() + self._splat_rain_gpu(d_display) + + # Advance volumetric cloud animation time + if self._clouds_enabled and self._volumetric_clouds_enabled: + self._cloud_time += 0.05 + # Sync: splat kernels run on stream 0, readback on non-blocking stream - if self._wind_enabled or self._hydro_enabled: + if _any_particles: sync_event = self._wind_done_event or self._hydro_done_event if sync_event is None: sync_event = cp.cuda.Event() @@ -8189,6 +8840,7 @@ def _render_help_text(self): ("DATA LAYERS", [ ("Shift+F", "FIRMS fire (7d)"), ("Shift+W", "Toggle wind"), + ("Shift+N", "Toggle clouds + rain"), ("Shift+Y", "Toggle hydro flow"), ]), ("RENDERING", [ @@ -8470,6 +9122,9 @@ def run(self, start_position: Optional[Tuple[float, float, float]] = None, _saved_termios = termios.tcgetattr(sys.stdin.fileno()) except (ImportError, OSError, ValueError): pass + except Exception: + # termios.error doesn't inherit from OSError on all platforms + pass # --- GLFW window creation --- # Force X11 on GLFW 3.4+ to get window decorations (title bar with diff --git a/rtxpy/viewer/cloud.py b/rtxpy/viewer/cloud.py new file mode 100644 index 0000000..f70c14e --- /dev/null +++ b/rtxpy/viewer/cloud.py @@ -0,0 +1,45 @@ +"""Cloud particle state for the interactive viewer.""" + + +class CloudState: + """Holds cloud particle simulation state. + + Cloud particles float above the terrain at a fixed altitude, + drift with the wind field, and render as soft white puffs with + density proportional to the interpolated cloud_cover field. + """ + + __slots__ = ( + 'clouds_enabled', + 'cloud_cover_grid', # (H, W) interpolated 0-1 cloud fraction + 'cloud_particles', # (N, 2) — row, col in full-res pixel coords + 'cloud_sizes', # (N,) — world-space radius per particle + 'cloud_alphas', # (N,) — base alpha per particle + 'cloud_ages', # (N,) + 'cloud_lifetimes', # (N,) + 'cloud_n_particles', + 'cloud_max_age', + 'cloud_altitude', # world-space Z for the cloud layer + 'cloud_min_depth', # min camera distance for rendering + 'cloud_terrain_np', # cached CPU terrain for projection + 'volumetric_clouds_enabled', # bool: volumetric ray-marched clouds + 'cloud_thickness', # world-space thickness of cloud slab + 'cloud_time', # animation time counter for wind drift + ) + + def __init__(self): + self.clouds_enabled = False + self.cloud_cover_grid = None + self.cloud_particles = None + self.cloud_sizes = None + self.cloud_alphas = None + self.cloud_ages = None + self.cloud_lifetimes = None + self.cloud_n_particles = 4000 + self.cloud_max_age = 300 + self.cloud_altitude = 0.0 + self.cloud_min_depth = 0.0 + self.cloud_terrain_np = None + self.volumetric_clouds_enabled = False + self.cloud_thickness = 0.0 + self.cloud_time = 0.0 diff --git a/rtxpy/viewer/keybindings.py b/rtxpy/viewer/keybindings.py index 217a9ce..4d0aefa 100644 --- a/rtxpy/viewer/keybindings.py +++ b/rtxpy/viewer/keybindings.py @@ -39,6 +39,7 @@ 'L': '_action_toggle_drone_glow', # Drone glow 'T': '_action_cycle_time', # Time-of-day 'Y': '_action_toggle_hydro', # Hydro flow particles + 'N': '_action_toggle_clouds', # Cloud layer } # Lowercase key bindings — checked after shift bindings