Atmosphere LUT parameterization improvements

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_4, fast_atan2}
 
 #import bevy_pbr::atmosphere::{
     types::Atmosphere,
@@ -38,7 +38,7 @@
 // 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
@@ -71,41 +71,34 @@ 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);
+    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);
+
+    // Latitude parameterization
+    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 mapping
+    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,12 +120,26 @@ 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);
+    let view_ray_dist = length(view_pos.xyz / view_pos.w) * settings.scene_units_to_m;
+    let t_max = settings.aerial_view_lut_max_distance;
+
+    // Special handling for first slice to avoid extrapolation
+    let delta_slice = t_max / f32(settings.aerial_view_lut_size.z);
+    if (view_ray_dist < delta_slice) {
+        let f = view_ray_dist / delta_slice;
+        let first_slice_uvw = vec3(uv, 0.5 / f32(settings.aerial_view_lut_size.z));
+        let sample = textureSampleLevel(aerial_view_lut, aerial_view_lut_sampler, first_slice_uvw, 0.0);
+        return vec4(sample.rgb * f, pow(sample.a, f));
+    }
+
+    // Offset by 0.5 slice to sample at the center of each slice
+    let normalized_depth = (view_ray_dist / t_max) * f32(settings.aerial_view_lut_size.z - 1u);
+    let w = (normalized_depth - 0.5) / f32(settings.aerial_view_lut_size.z);
+    let uvw = vec3(uv, clamp(w, 0.0, 1.0));
     return textureSampleLevel(aerial_view_lut, aerial_view_lut_sampler, uvw, 0.0);
 }
 
@@ -236,15 +243,18 @@ 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];
         let neg_LdotV = dot((*light).direction_to_light, ray_dir_ws);
-        let angle_to_sun = fast_acos(neg_LdotV);
+        let angle_to_sun = fast_acos_4(neg_LdotV);
         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_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.