Procedural Atmosphere PBR integration

crates/bevy_image/src/image_texture_conversion.rs

@@ -197,6 +197,22 @@ impl Image {
                 })
                 .map(DynamicImage::ImageRgba8)
             }
+            TextureFormat::Rgba16Float => {
+                use half::f16;
+                let pixel_count = (width * height) as usize;
+                let mut rgba32f_data = Vec::<f32>::with_capacity(pixel_count * 4);
+                for rgba16f in data.chunks_exact(8) {
+                    let r = f16::from_bits(u16::from_le_bytes([rgba16f[0], rgba16f[1]])).to_f32();
+                    let g = f16::from_bits(u16::from_le_bytes([rgba16f[2], rgba16f[3]])).to_f32();
+                    let b = f16::from_bits(u16::from_le_bytes([rgba16f[4], rgba16f[5]])).to_f32();
+                    let a = f16::from_bits(u16::from_le_bytes([rgba16f[6], rgba16f[7]])).to_f32();
+                    rgba32f_data.push(r);
+                    rgba32f_data.push(g);
+                    rgba32f_data.push(b);
+                    rgba32f_data.push(a);
+                }
+                ImageBuffer::from_raw(width, height, rgba32f_data).map(DynamicImage::ImageRgba32F)
+            }
             // Throw and error if conversion isn't supported
             texture_format => return Err(IntoDynamicImageError::UnsupportedFormat(texture_format)),
         }

crates/bevy_pbr/src/atmosphere/aerial_view_lut.wgsl

@@ -2,17 +2,16 @@
     mesh_view_types::{Lights, DirectionalLight},
     atmosphere::{
         types::{Atmosphere, AtmosphereSettings},
-        bindings::{atmosphere, settings, view, lights, aerial_view_lut_out},
+        bindings::settings,
         functions::{
             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, MIDPOINT_RATIO
+            uv_to_ray_direction, MIDPOINT_RATIO, get_view_position
         },
     }
 }
 
-
 @group(0) @binding(13) var aerial_view_lut_out: texture_storage_3d<rgba16float, write>;
 
 @compute
@@ -22,8 +21,8 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
 
     let uv = (vec2<f32>(idx.xy) + 0.5) / vec2<f32>(settings.aerial_view_lut_size.xy);
     let ray_dir = uv_to_ray_direction(uv);
-    let r = view_radius();
-    let mu = ray_dir.y;
+    let world_pos = get_view_position();
+    let r = length(world_pos);
     let t_max = settings.aerial_view_lut_max_distance;
 
     var prev_t = 0.0;
@@ -36,15 +35,17 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
             let dt = (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.xyz);
+            // Calculate position and up vector at the current sample point
+            let sample_pos = world_pos + ray_dir.xyz * t_i;
+            let local_r = length(sample_pos);
+            let local_up = normalize(sample_pos);
 
             let local_atmosphere = sample_atmosphere(local_r);
             let sample_optical_depth = local_atmosphere.extinction * dt;
             let sample_transmittance = exp(-sample_optical_depth);
 
             // evaluate one segment of the integral
-            var inscattering = sample_local_inscattering(local_atmosphere, ray_dir.xyz, local_r, local_up);
+            var inscattering = sample_local_inscattering(local_atmosphere, ray_dir.xyz, sample_pos, true);
 
             // Analytical integration of the single scattering term in the radiance transfer equation
             let s_int = (inscattering - inscattering * sample_transmittance) / local_atmosphere.extinction;

crates/bevy_pbr/src/atmosphere/bindings.wgsl

@@ -20,3 +20,8 @@
 @group(0) @binding(10) var sky_view_lut_sampler: sampler;
 @group(0) @binding(11) var aerial_view_lut: texture_3d<f32>;
 @group(0) @binding(12) var aerial_view_lut_sampler: sampler;
+@group(0) @binding(14) var directional_shadow_texture: texture_depth_2d_array;
+@group(0) @binding(15) var directional_shadow_sampler: sampler_comparison;
+@group(0) @binding(16) var blue_noise_texture: texture_2d<f32>;
+@group(0) @binding(17) var blue_noise_sampler: sampler;
+@group(0) @binding(18) var<uniform> probe_transform_buffer: mat4x4<f32>;

crates/bevy_pbr/src/atmosphere/environment.wgsl

@@ -0,0 +1,74 @@
+#import bevy_render::maths::{PI}
+#import bevy_pbr::atmosphere::{
+    types::{Atmosphere, AtmosphereSettings},
+    bindings::{atmosphere, settings, probe_transform_buffer},
+    functions::{max_atmosphere_distance, raymarch_atmosphere}
+}
+
+@group(0) @binding(13) var output: texture_storage_2d_array<rgba16float, write>;
+
+// Convert from cubemap face and UV to direction vector
+fn face_uv_to_direction(face: u32, uv: vec2<f32>) -> vec3<f32> {
+    // Convert UV from [0,1] to [-1,1]
+    let coords = 2.0 * uv - 1.0;
+
+    // Generate direction based on face
+    var dir: vec3<f32>;
+    switch face {
+        case 0u: { // +X
+            dir = vec3<f32>(1.0, -coords.y, coords.x);
+        }
+        case 1u: { // -X
+            dir = vec3<f32>(-1.0, -coords.y, -coords.x);
+        }
+        case 2u: { // +Y
+            dir = vec3<f32>(coords.x, 1.0, -coords.y);
+        }
+        case 3u: { // -Y
+            dir = vec3<f32>(coords.x, -1.0, coords.y);
+        }
+        case 4u: { // +Z
+            dir = vec3<f32>(coords.x, -coords.y, -1.0);
+        }
+        default: { // -Z (5)
+            dir = vec3<f32>(-coords.x, -coords.y, 1.0);
+        }
+    }
+
+    return normalize(dir);
+}
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let dimensions = textureDimensions(output);
+    let slice_index = global_id.z;
+
+    if (global_id.x >= dimensions.x || global_id.y >= dimensions.y || slice_index >= 6u) {
+        return;
+    }
+
+    // Calculate normalized UV coordinates for this pixel
+    let uv = vec2<f32>(
+        (f32(global_id.x) + 0.5) / f32(dimensions.x),
+        (f32(global_id.y) + 0.5) / f32(dimensions.y)
+    );
+
+    var world_pos = probe_transform_buffer[3].xyz;
+
+    // offset by the origin point of the atmosphere scene
+    world_pos += atmosphere.origin;
+
+    let r = length(world_pos);
+
+    let ray_dir_ws = face_uv_to_direction(slice_index, uv);
+    let up = normalize(world_pos);
+    let mu = dot(ray_dir_ws, up);
+    let raymarch_steps = 16.0;
+    let t_max = max_atmosphere_distance(r, mu);
+    let sample_count = mix(1.0, raymarch_steps, clamp(t_max * 0.01, 0.0, 1.0));
+    let result = raymarch_atmosphere(world_pos, ray_dir_ws, t_max, sample_count, uv, false, true, false);
+    let inscattering = result.inscattering;
+    let color = vec4<f32>(inscattering, 1.0);
+
+    textureStore(output, vec2<i32>(global_id.xy), i32(slice_index), color);
+}