move to cubemap

crates/bevy_pbr/src/atmosphere/aerial_view_lut.wgsl

@@ -6,20 +6,19 @@
         functions::{
             sample_transmittance_lut, sample_atmosphere, rayleigh, henyey_greenstein,
             sample_multiscattering_lut, AtmosphereSample, sample_local_inscattering,
-            get_local_r, get_local_up, view_radius, M_TO_KM, uv_to_ndc, position_ndc_to_world, depth_ndc_to_view_z
+            get_local_r, get_local_up, view_radius, uv_to_ndc, position_ndc_to_world, depth_ndc_to_view_z
         },
         bruneton_functions::{distance_to_top_atmosphere_boundary, distance_to_bottom_atmosphere_boundary,ray_intersects_ground}
     }
 }
 
 
-@group(0) @binding(13) var aerial_view_lut_out: texture_storage_3d<rgba16float, write>;
+@group(0) @binding(12) var aerial_view_lut_out: texture_storage_3d<rgba16float, write>;
 
 @compute
 @workgroup_size(16, 16, 1) //TODO: this approach makes it so closer slices get fewer samples. But we also expect those to have less scattering. So win/win?
 fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
     if any(idx.xy > settings.aerial_view_lut_size.xy) { return; }
-    var optical_depth: vec3<f32> = vec3(0.0);
 
     let uv = (vec2<f32>(idx.xy) + 0.5) / vec2<f32>(settings.aerial_view_lut_size.xy);
     let ray_dir = uv_to_ray_direction(uv); //TODO: negate for lighting calcs?
@@ -28,6 +27,7 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
 
     var prev_t = 0.0;
     var total_inscattering = vec3(0.0);
+    var optical_depth = vec3(0.0);
     for (var slice_i: i32 = i32(settings.aerial_view_lut_size.z - 1); slice_i >= 0; slice_i--) { //reversed loop to iterate raw depth values near->far 
         var sum_transmittance = 0.0;
         for (var step_i: i32 = i32(settings.aerial_view_lut_samples - 1); step_i >= 0; step_i--) { //same here
@@ -50,11 +50,9 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
             var local_inscattering = sample_local_inscattering(local_atmosphere, transmittance_to_sample, ray_dir.xyz, local_r, local_up);
             total_inscattering += local_inscattering * step_length;
             sum_transmittance += transmittance_to_sample.r + transmittance_to_sample.g + transmittance_to_sample.b;
-
-            let mean_transmittance = sum_transmittance / (f32(settings.aerial_view_lut_samples) * 3.0);
-            textureStore(aerial_view_lut_out, vec3(vec2<i32>(idx.xy), slice_i), vec4(total_inscattering, mean_transmittance));
-            //textureStore(aerial_view_lut_out, vec3(vec2<i32>(idx.xy), slice_i), vec4(optical_depth, step_length));
         }
+        let mean_transmittance = sum_transmittance / (f32(settings.aerial_view_lut_samples) * 3.0);
+        textureStore(aerial_view_lut_out, vec3(vec2<i32>(idx.xy), slice_i), vec4(total_inscattering, mean_transmittance));
     }
 }
 

crates/bevy_pbr/src/atmosphere/bindings.wgsl

@@ -15,7 +15,7 @@
 @group(0) @binding(5) var transmittance_lut_sampler: sampler;
 @group(0) @binding(6) var multiscattering_lut: texture_2d<f32>;
 @group(0) @binding(7) var multiscattering_lut_sampler: sampler;
-@group(0) @binding(8) var sky_view_lut: texture_2d<f32>;
+@group(0) @binding(8) var sky_view_lut: texture_cube<f32>;
 @group(0) @binding(9) var sky_view_lut_sampler: sampler;
 @group(0) @binding(10) var aerial_view_lut: texture_3d<f32>;
 @group(0) @binding(11) var aerial_view_lut_sampler: sampler;

crates/bevy_pbr/src/atmosphere/functions.wgsl

@@ -32,17 +32,15 @@ fn multiscattering_lut_uv_to_r_mu(uv: vec2<f32>) -> vec2<f32> {
     return vec2(r, mu);
 }
 
-fn sky_view_lut_lat_long_to_uv(lat: f32, long: f32) -> vec2<f32> {
-    let u = long * FRAC_PI + 0.5;
-    let v = sqrt(2 * abs(lat) * FRAC_PI) * sign(lat) * 0.5 + 0.5;
-    return vec2(u, v);
+fn sky_view_lut_squash_ray_dir(ray_dir_vs: vec3<f32>) -> vec3<f32> {
+    let new_y = sqrt(abs(ray_dir_vs.y)) * sign(ray_dir_vs.y);
+    return normalize(vec3(ray_dir_vs.x, new_y, ray_dir_vs.z));
 }
 
-fn sky_view_lut_uv_to_lat_long(uv: vec2<f32>) -> vec2<f32> {
-    let long = (uv.x - 0.5) * TAU;
-    let v_minus_half = uv.y - 0.5;
-    let lat = TAU * (v_minus_half * v_minus_half) * sign(v_minus_half);
-    return vec2(lat, long);
+fn sky_view_lut_unsquash_ray_dir(ray_dir_vs: vec3<f32>) -> vec3<f32> {
+    let abs_y = abs(ray_dir_vs.y);
+    let new_y = abs_y * abs_y * sign(ray_dir_vs.y);
+    return normalize(vec3(ray_dir_vs.x, new_y, ray_dir_vs.z));
 }
 
 // LUT SAMPLING
@@ -57,12 +55,9 @@ fn sample_multiscattering_lut(r: f32, mu: f32) -> vec3<f32> {
     return textureSampleLevel(multiscattering_lut, multiscattering_lut_sampler, uv, 0.0).rgb;
 }
 
-fn sample_sky_view_lut(ray_dir: vec3<f32>) -> vec3<f32> {
-    let lat_long = ray_dir_to_lat_long(ray_dir);
-    let uv = sky_view_lut_lat_long_to_uv(lat_long.x, lat_long.y);
-    let long = fract(abs(lat_long.y));
-    //return vec3(long, long, long);
-    return textureSampleLevel(sky_view_lut, sky_view_lut_sampler, uv, 0.0).rgb;
+fn sample_sky_view_lut(ray_dir_vs: vec3<f32>) -> vec3<f32> {
+    let ray_dir_vs_squashed = sky_view_lut_squash_ray_dir(ray_dir_vs);
+    return textureSampleLevel(sky_view_lut, sky_view_lut_sampler, ray_dir_vs_squashed, 0.0).rgb;
 }
 
 //RGB channels: total inscattered light along the camera ray to the current sample.
@@ -83,7 +78,6 @@ fn henyey_greenstein(neg_LdotV: f32) -> f32 {
     return FRAC_4_PI * (1.0 - g * g) / (denom * sqrt(denom));
 }
 
-
 // ATMOSPHERE SAMPLING
 
 struct AtmosphereSample {
@@ -144,6 +138,7 @@ fn sample_local_inscattering(local_atmosphere: AtmosphereSample, transmittance_t
 
 //TODO: make pr to specify light angular size on struct itself
 const SUN_ANGULAR_SIZE: f32 = 0.00436332; //angular radius of sun in radians
+//const SUN_ANGULAR_SIZE: f32 = 0.1;
 
 fn sample_sun_disk(ray_dir: vec3<f32>, transmittance: vec3<f32>) -> vec3<f32> {
     var sun_contribution = vec3(0.0);
@@ -221,11 +216,10 @@ fn uv_to_ray_direction(uv: vec2<f32>) -> vec4<f32> {
     return vec4(normalize(ray_direction), -view_ray_direction.z);
 }
 
-fn ray_dir_to_lat_long(ray_dir: vec3<f32>) -> vec2<f32> {
-    let view_dir = -view.world_from_view[2].xyz;
-    let lat = asin(ray_dir.y);
-    let long = atan2(view_dir.z, view_dir.x) - atan2(ray_dir.z, ray_dir.x); //TODO: explain
-    return vec2(lat, long);
+/// Convert a view space direction to world space
+fn direction_view_to_world(view_dir: vec3<f32>) -> vec3<f32> {
+    let world_dir = view.world_from_view * vec4(view_dir, 0.0);
+    return world_dir.xyz;
 }
 
 /// Convert ndc depth to linear view z. 

crates/bevy_pbr/src/atmosphere/mod.rs

@@ -178,25 +178,25 @@ pub struct Atmosphere {
     /// Radius of the planet
     ///
     /// units: km
-    bottom_radius: f32,
+    pub bottom_radius: f32,
 
     /// Radius at which we consider the atmosphere to 'end' for out calculations (from center of planet)
     /// units: km
-    top_radius: f32,
+    pub top_radius: f32,
 
-    ground_albedo: Vec3, //used for estimating multiscattering
+    pub ground_albedo: Vec3, //used for estimating multiscattering
 
-    rayleigh_density_exp_scale: f32,
-    rayleigh_scattering: Vec3,
+    pub rayleigh_density_exp_scale: f32,
+    pub rayleigh_scattering: Vec3,
 
-    mie_density_exp_scale: f32,
-    mie_scattering: f32, //units: km^-1
-    mie_absorption: f32, //units: km^-1
-    mie_asymmetry: f32,  //the "asymmetry" value of the phase function, unitless. Domain: (-1, 1)
+    pub mie_density_exp_scale: f32,
+    pub mie_scattering: f32, //units: km^-1
+    pub mie_absorption: f32, //units: km^-1
+    pub mie_asymmetry: f32, //the "asymmetry" value of the phase function, unitless. Domain: (-1, 1)
 
-    ozone_layer_center_altitude: f32, //units: km
-    ozone_layer_half_width: f32,      //units: km
-    ozone_absorption: Vec3,           //ozone absorption. units: km^-1
+    pub ozone_layer_center_altitude: f32, //units: km
+    pub ozone_layer_half_width: f32,      //units: km
+    pub ozone_absorption: Vec3,           //ozone absorption. units: km^-1
 }
 
 impl Atmosphere {
@@ -240,10 +240,10 @@ impl ExtractComponent for Atmosphere {
 pub struct AtmosphereSettings {
     pub transmittance_lut_size: UVec2,
     pub multiscattering_lut_size: UVec2,
-    pub sky_view_lut_size: UVec2,
+    pub aerial_view_lut_size: UVec3,
+    pub sky_view_lut_size: u32,
     pub multiscattering_lut_dirs: u32,
     pub transmittance_lut_samples: u32,
-    pub aerial_view_lut_size: UVec3,
     pub multiscattering_lut_samples: u32,
     pub sky_view_lut_samples: u32,
     pub aerial_view_lut_samples: u32,
@@ -258,7 +258,7 @@ impl Default for AtmosphereSettings {
             multiscattering_lut_size: UVec2::new(32, 32),
             multiscattering_lut_dirs: 64,
             multiscattering_lut_samples: 20,
-            sky_view_lut_size: UVec2::new(200, 100),
+            sky_view_lut_size: 64,
             sky_view_lut_samples: 30,
             aerial_view_lut_size: UVec3::new(32, 32, 32),
             aerial_view_lut_samples: 10,

crates/bevy_pbr/src/atmosphere/node.rs

@@ -62,7 +62,7 @@ impl ViewNode for AtmosphereLutsNode {
         ) = (
             pipeline_cache.get_render_pipeline(pipelines.transmittance_lut),
             pipeline_cache.get_compute_pipeline(pipelines.multiscattering_lut),
-            pipeline_cache.get_render_pipeline(pipelines.sky_view_lut),
+            pipeline_cache.get_compute_pipeline(pipelines.sky_view_lut),
             pipeline_cache.get_compute_pipeline(pipelines.aerial_view_lut),
         )
         else {
@@ -127,16 +127,14 @@ impl ViewNode for AtmosphereLutsNode {
             multiscattering_lut_pass.dispatch_workgroups(workgroups_x, workgroups_y, 1);
         }
 
+        const SKY_VIEW_WORKGROUP_SIZE: u32 = 16;
+        let workgroups_x = settings.sky_view_lut_size.div_ceil(SKY_VIEW_WORKGROUP_SIZE);
+        let workgroups_y = settings.sky_view_lut_size.div_ceil(SKY_VIEW_WORKGROUP_SIZE);
+
         {
-            let mut sky_view_lut_pass = commands.begin_render_pass(&RenderPassDescriptor {
+            let mut sky_view_lut_pass = commands.begin_compute_pass(&ComputePassDescriptor {
                 label: Some("sky_view_lut_pass"),
-                color_attachments: &[Some(RenderPassColorAttachment {
-                    view: &textures.sky_view_lut.default_view,
-                    resolve_target: None,
-                    ops: Operations::default(),
-                })],
-                depth_stencil_attachment: None,
-                ..Default::default()
+                timestamp_writes: None,
             });
             sky_view_lut_pass.set_pipeline(sky_view_lut_pipeline);
             sky_view_lut_pass.set_bind_group(
@@ -149,7 +147,7 @@ impl ViewNode for AtmosphereLutsNode {
                     lights_uniforms_offset.offset,
                 ],
             );
-            sky_view_lut_pass.draw(0..3, 0..1);
+            sky_view_lut_pass.dispatch_workgroups(workgroups_x, workgroups_y, 6);
         }
 
         {

crates/bevy_pbr/src/atmosphere/render_sky.wgsl

@@ -1,7 +1,7 @@
 #import bevy_pbr::atmosphere::{
     types::{Atmosphere, AtmosphereSettings},
     bindings::{atmosphere, view},
-    functions::{sample_transmittance_lut, sample_sky_view_lut, uv_to_ray_direction, uv_to_ndc, sample_aerial_view_lut, view_radius, sample_sun_disk},
+    functions::{sample_transmittance_lut, sample_sky_view_lut, direction_view_to_world, uv_to_ndc, sample_aerial_view_lut, view_radius, sample_sun_disk},
 };
 #import bevy_render::view::View;
 
@@ -13,17 +13,42 @@
 fn main(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
     let depth = textureLoad(depth_texture, vec2<i32>(in.position.xy), 0);
     if depth == 0.0 {
-        let ray_dir = uv_to_ray_direction(in.uv).xyz;
+        let view_ray_dir = uv_to_ray_direction(in.uv).xyz;
+        let world_ray_dir = direction_view_to_world(view_ray_dir);
         let r = view_radius();
-        let mu = ray_dir.y;
-        let sky_view = sample_sky_view_lut(ray_dir);
+        let mu = world_ray_dir.y;
+        let sky_view = sample_sky_view_lut(view_ray_dir);
         let transmittance = sample_transmittance_lut(r, mu);
-        let sun_disk = sample_sun_disk(ray_dir, transmittance);
-        return vec4(sky_view + sun_disk, 1.0);
+        let sun_disk = sample_sun_disk(world_ray_dir, transmittance);
+        return vec4(sky_view + sun_disk, (transmittance.r + transmittance.g + transmittance.b) / 3.0);
     } else {
         let ndc_xy = uv_to_ndc(in.uv);
         let ndc = vec3(ndc_xy, depth);
         let inscattering = sample_aerial_view_lut(ndc);
         return inscattering;
     }
 }
+
+//Modified from skybox.wgsl. For this pass we don't need to apply a separate sky transform or consider camera viewport.
+//w component is the cosine of the view direction with the view forward vector, to correct step distance at the edges of the viewport
+fn uv_to_ray_direction(uv: vec2<f32>) -> vec4<f32> {
+    // Using world positions of the fragment and camera to calculate a ray direction
+    // breaks down at large translations. This code only needs to know the ray direction.
+    // The ray direction is along the direction from the camera to the fragment position.
+    // In view space, the camera is at the origin, so the view space ray direction is
+    // along the direction of the fragment position - (0,0,0) which is just the
+    // fragment position.
+    // Use the position on the near clipping plane to avoid -inf world position
+    // because the far plane of an infinite reverse projection is at infinity.
+    let view_position_homogeneous = view.view_from_clip * vec4(
+        uv_to_ndc(uv),
+        1.0,
+        1.0,
+    );
+
+    // Transforming the view space ray direction by the skybox transform matrix, it is 
+    // equivalent to rotating the skybox itself.
+    let view_ray_direction = view_position_homogeneous.xyz / view_position_homogeneous.w; //TODO: remove this step and just use position_ndc_to_world? we didn't need to transform in view space
+
+    return vec4(normalize(view_ray_direction), -view_ray_direction.z);
+}