WIP: Atmosphere PBR support

crates/bevy_light/src/lib.rs

@@ -27,7 +27,10 @@ use cluster::{
 mod ambient_light;
 pub use ambient_light::AmbientLight;
 mod probe;
-pub use probe::{EnvironmentMapLight, GeneratedEnvironmentMapLight, IrradianceVolume, LightProbe};
+pub use probe::{
+    AtmosphereEnvironmentMapLight, EnvironmentMapLight, GeneratedEnvironmentMapLight,
+    IrradianceVolume, LightProbe,
+};
 mod volumetric;
 pub use volumetric::{FogVolume, VolumetricFog, VolumetricLight};
 pub mod cascade;

crates/bevy_light/src/probe.rs

@@ -2,7 +2,7 @@ use bevy_asset::Handle;
 use bevy_camera::visibility::Visibility;
 use bevy_ecs::prelude::*;
 use bevy_image::Image;
-use bevy_math::Quat;
+use bevy_math::{Quat, UVec2};
 use bevy_reflect::prelude::*;
 use bevy_transform::components::Transform;
 
@@ -140,6 +140,29 @@ impl Default for GeneratedEnvironmentMapLight {
     }
 }
 
+/// Generates a cubemap of the sky produced by the atmosphere shader.
+///
+/// Attach to a `LightProbe` to capture reflections inside its volume.
+#[derive(Component, Clone)]
+pub struct AtmosphereEnvironmentMapLight {
+    /// Luminance multiplier in cd/m².
+    pub intensity: f32,
+    /// Whether the diffuse contribution should affect meshes that already have lightmaps.
+    pub affects_lightmapped_mesh_diffuse: bool,
+    /// Cubemap resolution in pixels (must be a power-of-two).
+    pub size: UVec2,
+}
+
+impl Default for AtmosphereEnvironmentMapLight {
+    fn default() -> Self {
+        Self {
+            intensity: 1.0,
+            affects_lightmapped_mesh_diffuse: true,
+            size: UVec2::new(512, 512),
+        }
+    }
+}
+
 /// The component that defines an irradiance volume.
 ///
 /// See `bevy_pbr::irradiance_volume` for detailed information.

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, get_blue_noise
         },
     }
 }
 
-
 @group(0) @binding(13) var aerial_view_lut_out: texture_storage_3d<rgba16float, write>;
 
 @compute
@@ -21,9 +20,11 @@ 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);
+    // Apply blue-noise based jitter to the sampling offset
+    let sample_offset = (get_blue_noise(uv) - 0.5) * settings.jitter_strength + 0.5;
     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;
@@ -32,19 +33,21 @@ 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) + MIDPOINT_RATIO) / f32(settings.aerial_view_lut_samples))) / f32(settings.aerial_view_lut_size.z);
+            let t_i = t_max * (f32(slice_i) + ((f32(step_i) + sample_offset) / f32(settings.aerial_view_lut_samples))) / f32(settings.aerial_view_lut_size.z);
             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;
@@ -60,4 +63,4 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
         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, 0.0));
     }
-}
+}

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_array<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/bruneton_functions.wgsl

@@ -64,19 +64,19 @@
 // Assuming r between ground and top atmosphere boundary, and mu= cos(zenith_angle)
 // Chosen to increase precision near the ground and to work around a discontinuity at the horizon
 // See Bruneton and Neyret 2008, "Precomputed Atmospheric Scattering" section 4
-fn transmittance_lut_r_mu_to_uv(r: f32, mu: f32) -> vec2<f32> {
+fn transmittance_lut_r_mu_to_uv(atm: Atmosphere, r: f32, mu: f32) -> vec2<f32> {
   // Distance along a horizontal ray from the ground to the top atmosphere boundary
-    let H = sqrt(atmosphere.top_radius * atmosphere.top_radius - atmosphere.bottom_radius * atmosphere.bottom_radius);
+    let H = sqrt(atm.top_radius * atm.top_radius - atm.bottom_radius * atm.bottom_radius);
 
   // Distance from a point at height r to the horizon
   // ignore the case where r <= atmosphere.bottom_radius
-    let rho = sqrt(max(r * r - atmosphere.bottom_radius * atmosphere.bottom_radius, 0.0));
+    let rho = sqrt(max(r * r - atm.bottom_radius * atm.bottom_radius, 0.0));
 
   // Distance from a point at height r to the top atmosphere boundary at zenith angle mu
-    let d = distance_to_top_atmosphere_boundary(r, mu);
+    let d = distance_to_top_atmosphere_boundary(atm, r, mu);
 
   // Minimum and maximum distance to the top atmosphere boundary from a point at height r
-    let d_min = atmosphere.top_radius - r; // length of the ray straight up to the top atmosphere boundary
+    let d_min = atm.top_radius - r; // length of the ray straight up to the top atmosphere boundary
     let d_max = rho + H; // length of the ray to the top atmosphere boundary and grazing the horizon
 
     let u = (d - d_min) / (d_max - d_min);
@@ -121,9 +121,9 @@ fn transmittance_lut_uv_to_r_mu(uv: vec2<f32>) -> vec2<f32> {
 /// Center of sphere, c = [0,0,0]
 /// Radius of sphere, r = atmosphere.top_radius
 /// This function solves the quadratic equation for line-sphere intersection simplified under these assumptions
-fn distance_to_top_atmosphere_boundary(r: f32, mu: f32) -> f32 {
-  // ignore the case where r > atmosphere.top_radius
-    let positive_discriminant = max(r * r * (mu * mu - 1.0) + atmosphere.top_radius * atmosphere.top_radius, 0.0);
+fn distance_to_top_atmosphere_boundary(atm: Atmosphere, r: f32, mu: f32) -> f32 {
+  // ignore the case where r > atm.top_radius
+    let positive_discriminant = max(r * r * (mu * mu - 1.0) + atm.top_radius * atm.top_radius, 0.0);
     return max(-r * mu + sqrt(positive_discriminant), 0.0);
 }
 

crates/bevy_pbr/src/atmosphere/environment.wgsl

@@ -0,0 +1,46 @@
+#import bevy_render::maths::{PI}
+#import bevy_pbr::{
+    atmosphere::{
+        types::{Atmosphere, AtmosphereSettings},
+        bindings::{atmosphere, settings, probe_transform_buffer},
+        functions::{max_atmosphere_distance, raymarch_atmosphere},
+    },
+    utils::sample_cube_dir
+}
+
+@group(0) @binding(13) var output: texture_storage_2d_array<rgba16float, write>;
+
+@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 = sample_cube_dir(uv, slice_index);
+    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);
+}