WIP: Cloud scattering prototype

crates/bevy_pbr/src/atmosphere/aerial_view_lut.wgsl

@@ -16,7 +16,9 @@
 fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
     if any(idx.xy > settings.aerial_view_lut_size.xy) { return; }
 
-    let uv = (vec2<f32>(idx.xy) + 0.5) / vec2<f32>(settings.aerial_view_lut_size.xy);
+    // Use global invocation ID as pixel coordinates for jittering
+    let pixel_coords = vec2<f32>(idx.xy);
+    let uv = (pixel_coords + 0.5) / vec2<f32>(settings.aerial_view_lut_size.xy);
     let ray_dir = uv_to_ray_direction(uv);
     let world_pos = get_view_position();
 
@@ -45,7 +47,7 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
             let sample_transmittance = exp(-sample_optical_depth);
 
             // evaluate one segment of the integral
-            var inscattering = sample_local_inscattering(scattering, ray_dir, sample_pos);
+            var inscattering = sample_local_inscattering(scattering, ray_dir, sample_pos, pixel_coords);
 
             // Analytical integration of the single scattering term in the radiance transfer equation
             let s_int = (inscattering - inscattering * sample_transmittance) / max(extinction, MIN_EXTINCTION);

crates/bevy_pbr/src/atmosphere/bindings.wgsl

@@ -22,3 +22,15 @@
 @group(0) @binding(10) var sky_view_lut: texture_2d<f32>;
 @group(0) @binding(11) var aerial_view_lut: texture_3d<f32>;
 @group(0) @binding(12) var atmosphere_lut_sampler: sampler;
+
+// Cloud shadow map (front depth + extinction stats) generated by a compute pass.
+// Encodes:
+// - R: front depth (meters) from the light-volume near plane
+// - G: mean extinction (1/m)
+// - B: max optical depth (unitless)
+@group(0) @binding(17) var cloud_shadow_map: texture_2d<f32>;
+
+#ifdef CLOUDS_ENABLED
+// Spatio-temporal blue noise for per-frame stratification
+@group(0) @binding(19) var stbn_texture: texture_2d_array<f32>;
+#endif

crates/bevy_pbr/src/atmosphere/cloud_shadow_filter.wgsl

@@ -0,0 +1,136 @@
+#define_import_path bevy_pbr::atmosphere::cloud_shadow_filter
+
+#import bevy_pbr::atmosphere::bindings::{settings, atmosphere_lut_sampler}
+
+// Source cloud shadow map (rgba16f sampled as float).
+// Layout matches `bindings.wgsl`:
+// - R: front depth (km)
+// - G: mean extinction (1/m)
+// - B: max optical depth (unitless)
+@group(0) @binding(17) var cloud_shadow_src: texture_2d<f32>;
+
+// Destination ping-pong texture.
+@group(0) @binding(13) var cloud_shadow_dst: texture_storage_2d<rgba16float, write>;
+
+fn load_shadow(coord: vec2<i32>) -> vec3<f32> {
+    // Sample at texel centers using linear filtering (sampler clamps to edge).
+    let size = vec2<f32>(settings.cloud_shadow_map_size);
+    let uv = (vec2<f32>(coord) + vec2(0.5)) / size;
+    return textureSampleLevel(cloud_shadow_src, atmosphere_lut_sampler, uv, 0.0).rgb;
+}
+
+fn sort2(a: ptr<function, f32>, b: ptr<function, f32>) {
+    if ((*a) > (*b)) {
+        let t = (*a);
+        (*a) = (*b);
+        (*b) = t;
+    }
+}
+
+// Median of 9 values (3x3). Used to denoise front-depth without the “expanding rings”
+// you get from iterative min/max dilate/erode passes.
+fn median9(
+    v0: f32, v1: f32, v2: f32,
+    v3: f32, v4: f32, v5: f32,
+    v6: f32, v7: f32, v8: f32
+) -> f32 {
+    var a = v0; var b = v1; var c = v2;
+    var d = v3; var e = v4; var f = v5;
+    var g = v6; var h = v7; var i = v8;
+
+    sort2(&a, &b); sort2(&d, &e); sort2(&g, &h);
+    sort2(&b, &c); sort2(&e, &f); sort2(&h, &i);
+    sort2(&a, &b); sort2(&d, &e); sort2(&g, &h);
+
+    sort2(&a, &d); sort2(&d, &g);
+    sort2(&b, &e); sort2(&e, &h);
+    sort2(&c, &f); sort2(&f, &i);
+
+    sort2(&c, &e); sort2(&e, &g);
+    sort2(&c, &e); sort2(&e, &g);
+
+    return e;
+}
+
+// Spatial filter in transmittance space for max optical depth (Unreal-style idea):
+// - Filter visibility V = exp(-OD) with a small kernel
+// - Convert back: OD = -log(V)
+// Keep depth mostly unfiltered to avoid precision / convergence issues.
+@compute
+// Keep in sync with `dispatch_2d()` in `atmosphere/node.rs` (currently assumes 16x16 workgroups).
+@workgroup_size(16, 16, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let size = settings.cloud_shadow_map_size;
+    if (gid.x >= size.x || gid.y >= size.y) {
+        return;
+    }
+
+    let p = vec2<i32>(gid.xy);
+    let center = load_shadow(p);
+
+    // 3x3 tent-ish kernel weights:
+    // 1 2 1
+    // 2 4 2
+    // 1 2 1
+    // sum = 16
+    var sum_mean_ext = 0.0;
+    var sum_vis = 0.0;
+    var sum_w = 0.0;
+
+    for (var dy: i32 = -1; dy <= 1; dy += 1) {
+        for (var dx: i32 = -1; dx <= 1; dx += 1) {
+            let w = select(1.0, 2.0, (dx == 0) || (dy == 0));
+            let tap = load_shadow(p + vec2(dx, dy));
+            let mean_ext = tap.y;
+            let od = max(0.0, tap.z);
+            let vis = exp(-od);
+            sum_mean_ext += w * mean_ext;
+            sum_vis += w * vis;
+            sum_w += w;
+        }
+    }
+
+    let mean_ext_out = sum_mean_ext / max(1.0, sum_w);
+    let vis_out = sum_vis / max(1.0, sum_w);
+    // Avoid -inf / NaN from log(0).
+    let od_out = select(100.0, -log(max(vis_out, 1e-6)), vis_out > 0.0);
+
+    // Depth denoise:
+    // Use a conservative median filter on valid depth samples. This reduces noisy “salt-and-pepper”
+    // in the front-depth channel without the runaway growth you can get from repeated dilation.
+    //
+    // NOTE: This pass can run multiple iterations (ping-pong). To keep it stable over multiple
+    // passes, we only blend *partially* toward the median.
+    const DEPTH_INVALID: f32 = 1.0e9;
+    const DEPTH_DENOISE_ALPHA: f32 = 1.0;
+
+    let d00 = load_shadow(p + vec2(-1, -1)).x;
+    let d10 = load_shadow(p + vec2( 0, -1)).x;
+    let d20 = load_shadow(p + vec2( 1, -1)).x;
+    let d01 = load_shadow(p + vec2(-1,  0)).x;
+    let d11 = center.x;
+    let d21 = load_shadow(p + vec2( 1,  0)).x;
+    let d02 = load_shadow(p + vec2(-1,  1)).x;
+    let d12 = load_shadow(p + vec2( 0,  1)).x;
+    let d22 = load_shadow(p + vec2( 1,  1)).x;
+
+    // Treat non-positive depths as "invalid / no cloud" so they don't dominate the median.
+    let v00 = select(DEPTH_INVALID, d00, d00 > 0.0);
+    let v10 = select(DEPTH_INVALID, d10, d10 > 0.0);
+    let v20 = select(DEPTH_INVALID, d20, d20 > 0.0);
+    let v01 = select(DEPTH_INVALID, d01, d01 > 0.0);
+    let v11 = select(DEPTH_INVALID, d11, d11 > 0.0);
+    let v21 = select(DEPTH_INVALID, d21, d21 > 0.0);
+    let v02 = select(DEPTH_INVALID, d02, d02 > 0.0);
+    let v12 = select(DEPTH_INVALID, d12, d12 > 0.0);
+    let v22 = select(DEPTH_INVALID, d22, d22 > 0.0);
+
+    let med = median9(v00, v10, v20, v01, v11, v21, v02, v12, v22);
+    let depth_med = select(center.x, med, med < 1.0e8);
+    let depth_out = mix(center.x, depth_med, DEPTH_DENOISE_ALPHA);
+
+    let out_rgb = vec3(depth_out, max(0.0, mean_ext_out), max(0.0, od_out));
+    textureStore(cloud_shadow_dst, p, vec4(out_rgb, 0.0));
+}
+
+

crates/bevy_pbr/src/atmosphere/cloud_shadow_map.wgsl

@@ -0,0 +1,223 @@
+#define_import_path bevy_pbr::atmosphere::cloud_shadow_map
+
+// Computes an Unreal-style cloud shadow map storing:
+// R: front depth (kilometers) from the light-volume near plane (stored in km to fit fp16 range)
+// G: mean extinction (1/m)
+// B: max optical depth (unitless)
+//
+// This is later sampled at arbitrary world points to cheaply compute
+// transmittance: T = exp(-tau) where tau = mean_ext * (d_sample - d_front),
+// clamped to max optical depth.
+
+#import bevy_pbr::atmosphere::bindings::{atmosphere, settings, view, lights}
+#import bevy_pbr::atmosphere::clouds::{
+    cloud_layer_segment,
+    get_cloud_medium_density,
+    get_cloud_scattering_coeff,
+    get_cloud_absorption_coeff,
+}
+
+@group(0) @binding(19) var stbn_texture: texture_2d_array<f32>;
+@group(0) @binding(13) var out_cloud_shadow_map: texture_storage_2d<rgba16float, write>;
+
+const EPSILON_M: f32 = 1.0;
+
+// Simple deterministic hash utilities (no frame index required).
+// Used to jitter integration per shadow-map texel to mitigate banding.
+fn hash_u32(x_in: u32) -> u32 {
+    var x = x_in;
+    x ^= x >> 16u;
+    x *= 0x7feb352du;
+    x ^= x >> 15u;
+    x *= 0x846ca68bu;
+    x ^= x >> 16u;
+    return x;
+}
+
+fn hash_2d_to_01(p: vec2<u32>, salt: u32) -> f32 {
+    let h = hash_u32(p.x ^ (hash_u32(p.y) + 0x9e3779b9u) ^ salt);
+    // Convert to [0,1). 2^32 as f32 is representable.
+    return f32(h) * (1.0 / 4294967296.0);
+}
+
+fn safe_normalize(v: vec3<f32>) -> vec3<f32> {
+    let l2 = dot(v, v);
+    if (l2 <= 1e-12) {
+        return vec3(0.0, 1.0, 0.0);
+    }
+    return v * inverseSqrt(l2);
+}
+
+fn clamp_to_surface(position: vec3<f32>) -> vec3<f32> {
+    let min_radius = atmosphere.inner_radius + EPSILON_M;
+    let r = length(position);
+    if (r < min_radius) {
+        let up = safe_normalize(position);
+        return up * min_radius;
+    }
+    return position;
+}
+
+fn get_view_position() -> vec3<f32> {
+    // Matches `functions.wgsl::get_view_position()` to keep anchor consistent.
+    let atmo_pos = (atmosphere.world_to_atmosphere * vec4(view.world_position, 1.0)).xyz;
+    return clamp_to_surface(atmo_pos);
+}
+
+struct LightBasis {
+    x: vec3<f32>,
+    y: vec3<f32>,
+};
+
+fn build_light_basis(light_dir: vec3<f32>) -> LightBasis {
+    // Pick a vector not parallel to `light_dir`, then build an orthonormal basis.
+    // IMPORTANT: keep this selection stable; switching too early causes the basis to “flip”
+    // when the light gets moderately close to zenith, which looks like the shadow map drifts/rotates.
+    let world_up = vec3(0.0, 1.0, 0.0);
+    let a = select(world_up, vec3(1.0, 0.0, 0.0), abs(dot(light_dir, world_up)) > 0.999);
+    let x = safe_normalize(cross(a, light_dir));
+    let y = cross(light_dir, x);
+    return LightBasis(x, y);
+}
+
+fn snap_anchor_to_texel_grid(anchor: vec3<f32>, basis: LightBasis, extent: f32, size: vec2<u32>) -> vec3<f32> {
+    // Snap the anchor in the shadow-map XY plane to stabilize the map under camera motion.
+    // This mirrors Unreal's stabilized shadow-map anchor strategy.
+    let res = vec2<f32>(size);
+    let texel_size = (2.0 * extent) / max(res, vec2(1.0));
+
+    let ax = dot(anchor, basis.x);
+    let ay = dot(anchor, basis.y);
+
+    let snapped_ax = round(ax / texel_size.x) * texel_size.x;
+    let snapped_ay = round(ay / texel_size.y) * texel_size.y;
+
+    let dx = snapped_ax - ax;
+    let dy = snapped_ay - ay;
+    return anchor + basis.x * dx + basis.y * dy;
+}
+
+@compute
+@workgroup_size(16, 16, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let size = settings.cloud_shadow_map_size;
+    if (gid.x >= size.x || gid.y >= size.y) {
+        return;
+    }
+
+    // If there is no directional light, write "no clouds".
+    if (lights.n_directional_lights == 0u) {
+        let half_depth = settings.cloud_shadow_map_extent * 2.0;
+        let far_depth_km = (2.0 * half_depth) * 0.001;
+        textureStore(out_cloud_shadow_map, vec2<i32>(gid.xy), vec4(far_depth_km, 0.0, 0.0, 0.0));
+        return;
+    }
+
+    // IMPORTANT: trace direction must point from the light toward the scene (light -> surface),
+    // matching Unreal's convention. Bevy's `direction_to_light` points from the point toward the
+    // light (surface -> light), so we negate it here.
+    let trace_dir = safe_normalize(-lights.directional_lights[0].direction_to_light);
+    let basis = build_light_basis(trace_dir);
+
+    let extent = settings.cloud_shadow_map_extent;
+    // Unreal-style: derive depth range from extent (prevents pathological far near-plane positions at low sun).
+    let half_depth = extent * 2.0;
+    let strength = settings.cloud_shadow_map_strength;
+
+    // Map texel -> light-space XY in [-extent, extent].
+    let uv = (vec2<f32>(gid.xy) + vec2(0.5)) / vec2<f32>(size);
+    let xy = (uv - vec2(0.5)) * (2.0 * extent);
+
+    var anchor = get_view_position();
+    anchor = snap_anchor_to_texel_grid(anchor, basis, extent, size);
+    let ray_origin = anchor + basis.x * xy.x + basis.y * xy.y - trace_dir * half_depth;
+    let max_t = 2.0 * half_depth;
+    let max_t_km = max_t * 0.001;
+
+    // Intersect the ray with the cloud layer shell and clamp to our light volume depth.
+    let seg = cloud_layer_segment(ray_origin, trace_dir);
+    if (seg.z < 0.5) {
+        textureStore(out_cloud_shadow_map, vec2<i32>(gid.xy), vec4(max_t_km, 0.0, 0.0, 0.0));
+        return;
+    }
+
+    let t_start = max(0.0, seg.x);
+    let t_end = min(seg.y, max_t);
+    if (t_end <= t_start) {
+        textureStore(out_cloud_shadow_map, vec2<i32>(gid.xy), vec4(max_t_km, 0.0, 0.0, 0.0));
+        return;
+    }
+
+    // More samples when the sun is near the horizon (Unreal-style):
+    // the ray travels a lot longer through the cloud shell, so a fixed sample count can miss
+    // density entirely, producing sweeping “no shadow” bands as the sun moves.
+    //
+    // Unreal uses:
+    //   HorizonFactor = clamp(0.2 / abs(dot(PlanetUp, -LightDir)), 0, 1)
+    //   SampleCount *= lerp(1, HorizonMultiplier, HorizonFactor)
+    //
+    // Here we approximate PlanetUp using the anchor up vector (good enough for ground/near-ground views).
+    let base_n = max(1u, settings.cloud_shadow_map_samples);
+    let anchor_up = safe_normalize(anchor);
+    let mu = abs(dot(anchor_up, trace_dir));
+    let horizon_factor = clamp(0.2 / max(mu, 1e-3), 0.0, 1.0);
+    const HORIZON_MULT: f32 = 2.0;
+    let n = min(512u, max(1u, u32(ceil(f32(base_n) * mix(1.0, HORIZON_MULT, horizon_factor)))));
+    // let n = 24u;
+    let dt = (t_end - t_start) / f32(n);
+
+    let sigma_t_per_density = get_cloud_scattering_coeff() + get_cloud_absorption_coeff();
+
+    var front_depth = max_t;
+    var sum_ext = 0.0;
+    var count_ext = 0.0;
+    var max_optical_depth = 0.0;
+
+    // Avoid “phantom” front depths from tiny numerical densities, which can cast shadows
+    // in empty space on the sun-facing side of clouds.
+    // Units: optical depth is unitless.
+    const HIT_OPTICAL_DEPTH_THRESHOLD: f32 = 1e-3;
+
+    for (var i = 0u; i < n; i += 1u) {
+        // Jittered + stratified: one sample per stratum, with *per-step* jitter.
+        // Per-step jitter reduces “streaky” structure compared to using a single phase shift
+        // across all steps in a texel.
+        var j: f32;
+        let stbn_dims = textureDimensions(stbn_texture);
+        if (all(stbn_dims > vec2(1u))) {
+            let stbn_layers = textureNumLayers(stbn_texture);
+            let stbn_layer = i32(view.frame_count % u32(stbn_layers));
+            let stbn_noise = textureLoad(stbn_texture, vec2<i32>(gid.xy) % vec2<i32>(stbn_dims), stbn_layer, 0);
+            // j = fract(stbn_noise.r + f32(i) * 0.618033988749895);
+            j = fract(0.0 + f32(i) * 0.618033988749895);
+        } else {
+            j = hash_2d_to_01(gid.xy, 0x1000u + i);
+        }
+        let t = t_start + (f32(i) + j) * dt;
+        let p = ray_origin + trace_dir * t;
+        let r = length(p);
+        let density = get_cloud_medium_density(r, p);
+        if (density > 0.0) {
+            let extinction = density * sigma_t_per_density; // 1/m
+            let od = extinction * dt;
+            if (od > HIT_OPTICAL_DEPTH_THRESHOLD) {
+                if (front_depth == max_t) {
+                    front_depth = t;
+                }
+                sum_ext += extinction;
+                count_ext += 1.0;
+                max_optical_depth += od;
+            }
+        }
+    }
+
+    let mean_ext = select(0.0, sum_ext / count_ext, count_ext > 0.0);
+
+    // Apply user strength scaling (matches Unreal's "strength" behavior).
+    let out_mean_ext = mean_ext * strength;
+    let out_max_od = max_optical_depth * strength;
+
+    textureStore(out_cloud_shadow_map, vec2<i32>(gid.xy), vec4(front_depth * 0.001, out_mean_ext, out_max_od, 0.0));
+}
+
+