Atmosphere LUT parameterization improvements (bevyengine#17555)

# Objective

- Fix the atmosphere LUT parameterization in the aerial -view and
sky-view LUTs
- Correct the light accumulation according to a ray-marched reference
- Avoid negative values of the sun disk illuminance when the sun disk is
below the horizon

## Solution

- Adding a Newton's method iteration to `fast_sqrt` function
- Switched to using `fast_acos_4` for better precision of the sun angle
towards the horizon (view mu angle = 0)
- Simplified the function for mapping to and from the Sky View UV
coordinates by removing an if statement and correctly apply the method
proposed by the [Hillarie
paper](https://sebh.github.io/publications/egsr2020.pdf) detailed in
section 5.3 and 5.4.
- Replaced the `ray_dir_ws.y` term with a shadow factor in the
`sample_sun_illuminance` function that correctly approximates the sun
disk occluded by the earth from any view point

## Testing

- Ran the atmosphere and SSAO examples to make sure the shaders still
compile and run as expected.

---

## Showcase

<img width="1151" alt="showcase-img"
src="https://github.com/user-attachments/assets/de875533-42bd-41f9-9fd0-d7cc57d6e51c"
/>

---------

Co-authored-by: Emerson Coskey <[email protected]>

crates/bevy_pbr/src/atmosphere/aerial_view_lut.wgsl

@@ -7,7 +7,7 @@
             sample_transmittance_lut, sample_atmosphere, rayleigh, henyey_greenstein,
             sample_multiscattering_lut, AtmosphereSample, sample_local_inscattering,
             get_local_r, get_local_up, view_radius, uv_to_ndc, max_atmosphere_distance,
-            uv_to_ray_direction
+            uv_to_ray_direction, MIDPOINT_RATIO
         },
     }
 }
@@ -32,7 +32,7 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
 
     for (var slice_i: u32 = 0; slice_i < settings.aerial_view_lut_size.z; slice_i++) {
         for (var step_i: u32 = 0; step_i < settings.aerial_view_lut_samples; step_i++) {
-            let t_i = t_max * (f32(slice_i) + ((f32(step_i) + 0.5) / f32(settings.aerial_view_lut_samples))) / f32(settings.aerial_view_lut_size.z);
+            let t_i = t_max * (f32(slice_i) + ((f32(step_i) + MIDPOINT_RATIO) / f32(settings.aerial_view_lut_samples))) / f32(settings.aerial_view_lut_size.z);
             let dt = (t_i - prev_t);
             prev_t = t_i;
 
@@ -55,8 +55,12 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
                 break;
             }
         }
-        //We only have one channel to store transmittance, so we store the mean
+        // We only have one channel to store transmittance, so we store the mean
         let mean_transmittance = (throughput.r + throughput.g + throughput.b) / 3.0;
-        textureStore(aerial_view_lut_out, vec3(vec2<u32>(idx.xy), slice_i), vec4(total_inscattering, mean_transmittance));
+
+        // Store in log space to allow linear interpolation of exponential values between slices
+        let log_transmittance = -log(max(mean_transmittance, 1e-6)); // Avoid log(0)
+        let log_inscattering = log(max(total_inscattering, vec3(1e-6)));
+        textureStore(aerial_view_lut_out, vec3(vec2<u32>(idx.xy), slice_i), vec4(log_inscattering, log_transmittance));
     }
 }

crates/bevy_pbr/src/atmosphere/functions.wgsl

@@ -1,6 +1,6 @@
 #define_import_path bevy_pbr::atmosphere::functions
 
-#import bevy_render::maths::{PI, HALF_PI, PI_2, fast_acos, fast_atan2}
+#import bevy_render::maths::{PI, HALF_PI, PI_2, fast_acos, fast_acos_4, fast_atan2}
 
 #import bevy_pbr::atmosphere::{
     types::Atmosphere,
@@ -38,11 +38,17 @@
 // CONSTANTS
 
 const FRAC_PI: f32 = 0.3183098862; // 1 / π
-const FRAC_2_PI: f32 = 0.15915494309;
+const FRAC_2_PI: f32 = 0.15915494309;  // 1 / (2π)
 const FRAC_3_16_PI: f32 = 0.0596831036594607509; // 3 / (16π)
 const FRAC_4_PI: f32 = 0.07957747154594767; // 1 / (4π)
 const ROOT_2: f32 = 1.41421356; // √2
 
+// During raymarching, each segment is sampled at a single point. This constant determines
+// where in the segment that sample is taken (0.0 = start, 0.5 = middle, 1.0 = end).
+// We use 0.3 to sample closer to the start of each segment, which better approximates
+// the exponential falloff of atmospheric density.
+const MIDPOINT_RATIO: f32 = 0.3;
+
 // LUT UV PARAMATERIZATIONS
 
 fn unit_to_sub_uvs(val: vec2<f32>, resolution: vec2<f32>) -> vec2<f32> {
@@ -71,41 +77,38 @@ fn sky_view_lut_r_mu_azimuth_to_uv(r: f32, mu: f32, azimuth: f32) -> vec2<f32> {
 
     let v_horizon = sqrt(r * r - atmosphere.bottom_radius * atmosphere.bottom_radius);
     let cos_beta = v_horizon / r;
-    let beta = fast_acos(cos_beta);
+    // Using fast_acos_4 for better precision at small angles
+    // to avoid artifacts at the horizon
+    let beta = fast_acos_4(cos_beta);
     let horizon_zenith = PI - beta;
-    let view_zenith = fast_acos(mu);
-
-    var v: f32;
-    if !ray_intersects_ground(r, mu) {
-        let coord = sqrt(1.0 - view_zenith / horizon_zenith);
-        v = (1.0 - coord) * 0.5;
-    } else {
-        let coord = (view_zenith - horizon_zenith) / beta;
-        v = sqrt(coord) * 0.5 + 0.5;
-    }
+    let view_zenith = fast_acos_4(mu);
+
+    // Apply non-linear transformation to compress more texels 
+    // near the horizon where high-frequency details matter most
+    // l is latitude in [-π/2, π/2] and v is texture coordinate in [0,1]
+    let l = view_zenith - horizon_zenith;
+    let abs_l = abs(l);
+
+    let v = 0.5 + 0.5 * sign(l) * sqrt(abs_l / HALF_PI);
 
     return unit_to_sub_uvs(vec2(u, v), vec2<f32>(settings.sky_view_lut_size));
 }
 
 fn sky_view_lut_uv_to_zenith_azimuth(r: f32, uv: vec2<f32>) -> vec2<f32> {
-    let adj_uv = sub_uvs_to_unit(uv, vec2<f32>(settings.sky_view_lut_size));
+    let adj_uv = sub_uvs_to_unit(vec2(uv.x, 1.0 - uv.y), vec2<f32>(settings.sky_view_lut_size));
     let azimuth = (adj_uv.x - 0.5) * PI_2;
 
+    // Horizon parameters
     let v_horizon = sqrt(r * r - atmosphere.bottom_radius * atmosphere.bottom_radius);
     let cos_beta = v_horizon / r;
-    let beta = fast_acos(cos_beta);
+    let beta = fast_acos_4(cos_beta);
     let horizon_zenith = PI - beta;
 
-    var zenith: f32;
-    if adj_uv.y < 0.5 {
-        let coord = 1.0 - 2.0 * adj_uv.y;
-        zenith = horizon_zenith * (1.0 - coord * coord);
-    } else {
-        let coord = 2.0 * adj_uv.y - 1.0;
-        zenith = horizon_zenith + beta * coord * coord;
-    }
+    // Inverse of horizon-detail mapping to recover original latitude from texture coordinate
+    let t = abs(2.0 * (adj_uv.y - 0.5));
+    let l = sign(adj_uv.y - 0.5) * HALF_PI * t * t;
 
-    return vec2(zenith, azimuth);
+    return vec2(horizon_zenith - l, azimuth);
 }
 
 // LUT SAMPLING
@@ -127,13 +130,23 @@ fn sample_sky_view_lut(r: f32, ray_dir_as: vec3<f32>) -> vec3<f32> {
     return textureSampleLevel(sky_view_lut, sky_view_lut_sampler, uv, 0.0).rgb;
 }
 
-//RGB channels: total inscattered light along the camera ray to the current sample.
-//A channel: average transmittance across all wavelengths to the current sample.
+// RGB channels: total inscattered light along the camera ray to the current sample.
+// A channel: average transmittance across all wavelengths to the current sample.
 fn sample_aerial_view_lut(uv: vec2<f32>, depth: f32) -> vec4<f32> {
     let view_pos = view.view_from_clip * vec4(uv_to_ndc(uv), depth, 1.0);
     let dist = length(view_pos.xyz / view_pos.w) * settings.scene_units_to_m;
-    let uvw = vec3(uv, dist / settings.aerial_view_lut_max_distance);
-    return textureSampleLevel(aerial_view_lut, aerial_view_lut_sampler, uvw, 0.0);
+    let t_max = settings.aerial_view_lut_max_distance;
+    let num_slices = f32(settings.aerial_view_lut_size.z);
+    // Offset the W coordinate by -0.5 over the max distance in order to 
+    // align sampling position with slice boundaries, since each texel 
+    // stores the integral over its entire slice
+    let uvw = vec3(uv, saturate(dist / t_max - 0.5 / num_slices));
+    let sample = textureSampleLevel(aerial_view_lut, aerial_view_lut_sampler, uvw, 0.0);
+    // Treat the first slice specially since there is 0 scattering at the camera
+    let delta_slice = t_max / num_slices;
+    let fade = saturate(dist / delta_slice);
+    // Recover the values from log space
+    return exp(sample) * fade;
 }
 
 // PHASE FUNCTIONS
@@ -236,6 +249,9 @@ fn sample_local_inscattering(local_atmosphere: AtmosphereSample, ray_dir: vec3<f
 const SUN_ANGULAR_SIZE: f32 = 0.0174533; // angular diameter of sun in radians
 
 fn sample_sun_illuminance(ray_dir_ws: vec3<f32>, transmittance: vec3<f32>) -> vec3<f32> {
+    let r = view_radius();
+    let mu_view = ray_dir_ws.y;
+    let shadow_factor = f32(!ray_intersects_ground(r, mu_view));
     var sun_illuminance = vec3(0.0);
     for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
         let light = &lights.directional_lights[light_i];
@@ -244,7 +260,7 @@ fn sample_sun_illuminance(ray_dir_ws: vec3<f32>, transmittance: vec3<f32>) -> ve
         let pixel_size = fwidth(angle_to_sun);
         let factor = smoothstep(0.0, -pixel_size * ROOT_2, angle_to_sun - SUN_ANGULAR_SIZE * 0.5);
         let sun_solid_angle = (SUN_ANGULAR_SIZE * SUN_ANGULAR_SIZE) * 4.0 * FRAC_PI;
-        sun_illuminance += ((*light).color.rgb / sun_solid_angle) * factor * ray_dir_ws.y;
+        sun_illuminance += ((*light).color.rgb / sun_solid_angle) * factor * shadow_factor;
     }
     return sun_illuminance * transmittance * view.exposure;
 }

crates/bevy_pbr/src/atmosphere/sky_view_lut.wgsl

@@ -8,6 +8,7 @@
             sample_local_inscattering, get_local_r, view_radius,
             max_atmosphere_distance, direction_atmosphere_to_world,
             sky_view_lut_uv_to_zenith_azimuth, zenith_azimuth_to_ray_dir,
+            MIDPOINT_RATIO
         },
     }
 }
@@ -34,21 +35,14 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
     let mu = ray_dir_ws.y;
     let t_max = max_atmosphere_distance(r, mu);
 
-    // Raymarch with quadratic distribution
     let sample_count = mix(1.0, f32(settings.sky_view_lut_samples), clamp(t_max * 0.01, 0.0, 1.0));
-    let sample_count_floor = floor(sample_count);
-    let t_max_floor = t_max * sample_count_floor / sample_count;
     var total_inscattering = vec3(0.0);
     var throughput = vec3(1.0);
+    var prev_t = 0.0;
     for (var s = 0.0; s < sample_count; s += 1.0) {
-        // Use quadratic distribution like reference
-        var t0 = (s / sample_count_floor);
-        var t1 = ((s + 1.0) / sample_count_floor);
-        t0 = t0 * t0;
-        t1 = t1 * t1;
-        t1 = select(t_max_floor * t1, t_max, t1 > 1.0);
-        let t_i = t_max_floor * t0 + (t1 - t_max_floor * t0) * 0.3;
-        let dt_i = t1 - t_max_floor * t0;
+        let t_i = t_max * (s + MIDPOINT_RATIO) / sample_count;
+        let dt_i = (t_i - prev_t);
+        prev_t = t_i;
 
         let local_r = get_local_r(r, mu, t_i);
         let local_up = get_local_up(r, t_i, ray_dir_ws);

crates/bevy_pbr/src/atmosphere/transmittance_lut.wgsl

@@ -1,7 +1,7 @@
 #import bevy_pbr::atmosphere::{
     types::{Atmosphere, AtmosphereSettings},
     bindings::{settings, atmosphere},
-    functions::{AtmosphereSample, sample_atmosphere, get_local_r, max_atmosphere_distance},
+    functions::{AtmosphereSample, sample_atmosphere, get_local_r, max_atmosphere_distance, MIDPOINT_RATIO},
     bruneton_functions::{transmittance_lut_uv_to_r_mu, distance_to_bottom_atmosphere_boundary, distance_to_top_atmosphere_boundary},
 }
 
@@ -32,7 +32,7 @@ fn ray_optical_depth(r: f32, mu: f32, sample_count: u32) -> vec3<f32> {
     var prev_t = 0.0f;
 
     for (var i = 0u; i < sample_count; i++) {
-        let t_i = t_max * (f32(i) + 0.3f) / f32(sample_count);
+        let t_i = t_max * (f32(i) + MIDPOINT_RATIO) / f32(sample_count);
         let dt = t_i - prev_t;
         prev_t = t_i;
 

crates/bevy_render/src/maths.wgsl

@@ -105,7 +105,9 @@ fn project_onto(lhs: vec3<f32>, rhs: vec3<f32>) -> vec3<f32> {
 // accuracy can be sacrificed for greater sample count.
 
 fn fast_sqrt(x: f32) -> f32 {
-    return bitcast<f32>(0x1fbd1df5 + (bitcast<i32>(x) >> 1u));
+    let n = bitcast<f32>(0x1fbd1df5 + (bitcast<i32>(x) >> 1u));
+    // One Newton's method iteration for better precision
+    return 0.5 * (n + x / n);
 }
 
 // Slightly less accurate than fast_acos_4, but much simpler.