Generalized atmospheric scattering media

crates/bevy_pbr/src/atmosphere/aerial_view_lut.wgsl

@@ -4,16 +4,17 @@
         types::{Atmosphere, AtmosphereSettings},
         bindings::{atmosphere, settings, view, lights, aerial_view_lut_out},
         functions::{
-            sample_transmittance_lut, sample_atmosphere, rayleigh, henyey_greenstein,
+            sample_transmittance_lut, sample_density_lut, rayleigh, henyey_greenstein,
             sample_multiscattering_lut, AtmosphereSample, sample_local_inscattering,
             uv_to_ndc, max_atmosphere_distance, uv_to_ray_direction, 
-            MIDPOINT_RATIO, get_view_position
+            MIDPOINT_RATIO, get_view_position, MIN_EXTINCTION, ABSORPTION_DENSITY,
+            SCATTERING_DENSITY,
         },
     }
 }
 
 
-@group(0) @binding(13) var aerial_view_lut_out: texture_storage_3d<rgba16float, write>;
+@group(0) @binding(16) var aerial_view_lut_out: texture_storage_3d<rgba16float, write>;
 
 @compute
 @workgroup_size(16, 16, 1)
@@ -23,7 +24,7 @@ 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 world_pos = get_view_position();
-    
+
     let r = length(world_pos);
     let t_max = settings.aerial_view_lut_max_distance;
 
@@ -41,15 +42,18 @@ fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
             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 absorption = sample_density_lut(local_r, ABSORPTION_DENSITY);
+            let scattering = sample_density_lut(local_r, SCATTERING_DENSITY);
+            let extinction = absorption + scattering;
+
+            let sample_optical_depth = extinction * dt;
             let sample_transmittance = exp(-sample_optical_depth);
 
             // evaluate one segment of the integral
-            var inscattering = sample_local_inscattering(local_atmosphere, ray_dir, sample_pos);
+            var inscattering = sample_local_inscattering(scattering, ray_dir, sample_pos);
 
             // Analytical integration of the single scattering term in the radiance transfer equation
-            let s_int = (inscattering - inscattering * sample_transmittance) / local_atmosphere.extinction;
+            let s_int = (inscattering - inscattering * sample_transmittance) / max(extinction, MIN_EXTINCTION);
             total_inscattering += throughput * s_int;
 
             throughput *= sample_transmittance;

crates/bevy_pbr/src/atmosphere/bindings.wgsl

@@ -12,11 +12,16 @@
 @group(0) @binding(2) var<uniform> atmosphere_transforms: AtmosphereTransforms;
 @group(0) @binding(3) var<uniform> view: View;
 @group(0) @binding(4) var<uniform> lights: Lights;
-@group(0) @binding(5) var transmittance_lut: texture_2d<f32>;
-@group(0) @binding(6) var transmittance_lut_sampler: sampler;
-@group(0) @binding(7) var multiscattering_lut: texture_2d<f32>;
-@group(0) @binding(8) var multiscattering_lut_sampler: sampler;
-@group(0) @binding(9) var sky_view_lut: texture_2d<f32>;
-@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(5) var medium_density_lut: texture_2d<f32>;
+@group(0) @binding(6) var medium_scattering_lut: texture_2d<f32>;
+@group(0) @binding(7) var medium_sampler: sampler;
+
+@group(0) @binding(8) var transmittance_lut: texture_2d<f32>;
+@group(0) @binding(9) var transmittance_lut_sampler: sampler;
+@group(0) @binding(10) var multiscattering_lut: texture_2d<f32>;
+@group(0) @binding(11) var multiscattering_lut_sampler: sampler;
+@group(0) @binding(12) var sky_view_lut: texture_2d<f32>;
+@group(0) @binding(13) var sky_view_lut_sampler: sampler;
+@group(0) @binding(14) var aerial_view_lut: texture_3d<f32>;
+@group(0) @binding(15) var aerial_view_lut_sampler: sampler;

crates/bevy_pbr/src/atmosphere/environment.rs

@@ -1,9 +1,9 @@
 use crate::{
     resources::{
         AtmosphereSamplers, AtmosphereTextures, AtmosphereTransform, AtmosphereTransforms,
-        AtmosphereTransformsOffset,
+        AtmosphereTransformsOffset, GpuAtmosphere,
     },
-    GpuAtmosphereSettings, GpuLights, LightMeta, ViewLightsUniformOffset,
+    ExtractedAtmosphere, GpuAtmosphereSettings, GpuLights, LightMeta, ViewLightsUniformOffset,
 };
 use bevy_asset::{load_embedded_asset, AssetServer, Assets, Handle, RenderAssetUsages};
 use bevy_ecs::{
@@ -29,8 +29,6 @@ use bevy_render::{
 use bevy_utils::default;
 use tracing::warn;
 
-use super::Atmosphere;
-
 // Render world representation of an environment map light for the atmosphere
 #[derive(Component, ExtractComponent, Clone)]
 pub struct AtmosphereEnvironmentMap {
@@ -65,26 +63,33 @@ pub struct AtmosphereProbePipeline {
 pub fn init_atmosphere_probe_layout(mut commands: Commands) {
     let environment = BindGroupLayoutDescriptor::new(
         "environment_bind_group_layout",
-        &BindGroupLayoutEntries::sequential(
+        &BindGroupLayoutEntries::with_indices(
             ShaderStages::COMPUTE,
             (
-                uniform_buffer::<Atmosphere>(true),
-                uniform_buffer::<GpuAtmosphereSettings>(true),
-                uniform_buffer::<AtmosphereTransform>(true),
-                uniform_buffer::<ViewUniform>(true),
-                uniform_buffer::<GpuLights>(true),
-                texture_2d(TextureSampleType::Float { filterable: true }), //transmittance lut and sampler
-                sampler(SamplerBindingType::Filtering),
-                texture_2d(TextureSampleType::Float { filterable: true }), //multiscattering lut and sampler
-                sampler(SamplerBindingType::Filtering),
-                texture_2d(TextureSampleType::Float { filterable: true }), //sky view lut and sampler
-                sampler(SamplerBindingType::Filtering),
-                texture_3d(TextureSampleType::Float { filterable: true }), //aerial view lut ans sampler
-                sampler(SamplerBindingType::Filtering),
-                texture_storage_2d_array(
-                    // output 2D array texture
-                    TextureFormat::Rgba16Float,
-                    StorageTextureAccess::WriteOnly,
+                (0, uniform_buffer::<GpuAtmosphere>(true)),
+                (1, uniform_buffer::<GpuAtmosphereSettings>(true)),
+                (2, uniform_buffer::<AtmosphereTransform>(true)),
+                (3, uniform_buffer::<ViewUniform>(true)),
+                (4, uniform_buffer::<GpuLights>(true)),
+                //transmittance lut and sampler
+                (8, texture_2d(TextureSampleType::default())),
+                (9, sampler(SamplerBindingType::Filtering)),
+                //multiscattering lut and sampler
+                (10, texture_2d(TextureSampleType::default())),
+                (11, sampler(SamplerBindingType::Filtering)),
+                //sky view lut and sampler
+                (12, texture_2d(TextureSampleType::default())),
+                (13, sampler(SamplerBindingType::Filtering)),
+                //aerial view lut ans sampler
+                (14, texture_3d(TextureSampleType::default())),
+                (15, sampler(SamplerBindingType::Filtering)),
+                // output 2D array texture
+                (
+                    16,
+                    texture_storage_2d_array(
+                        TextureFormat::Rgba16Float,
+                        StorageTextureAccess::WriteOnly,
+                    ),
                 ),
             ),
         ),
@@ -101,7 +106,7 @@ pub(super) fn prepare_atmosphere_probe_bind_groups(
     view_uniforms: Res<ViewUniforms>,
     lights_uniforms: Res<LightMeta>,
     atmosphere_transforms: Res<AtmosphereTransforms>,
-    atmosphere_uniforms: Res<ComponentUniforms<Atmosphere>>,
+    atmosphere_uniforms: Res<ComponentUniforms<GpuAtmosphere>>,
     settings_uniforms: Res<ComponentUniforms<GpuAtmosphereSettings>>,
     pipeline_cache: Res<PipelineCache>,
     mut commands: Commands,
@@ -110,21 +115,21 @@ pub(super) fn prepare_atmosphere_probe_bind_groups(
         let environment = render_device.create_bind_group(
             "environment_bind_group",
             &pipeline_cache.get_bind_group_layout(&layouts.environment),
-            &BindGroupEntries::sequential((
-                atmosphere_uniforms.binding().unwrap(),
-                settings_uniforms.binding().unwrap(),
-                atmosphere_transforms.uniforms().binding().unwrap(),
-                view_uniforms.uniforms.binding().unwrap(),
-                lights_uniforms.view_gpu_lights.binding().unwrap(),
-                &textures.transmittance_lut.default_view,
-                &samplers.transmittance_lut,
-                &textures.multiscattering_lut.default_view,
-                &samplers.multiscattering_lut,
-                &textures.sky_view_lut.default_view,
-                &samplers.sky_view_lut,
-                &textures.aerial_view_lut.default_view,
-                &samplers.aerial_view_lut,
-                &textures.environment,
+            &BindGroupEntries::with_indices((
+                (0, atmosphere_uniforms.binding().unwrap()),
+                (1, settings_uniforms.binding().unwrap()),
+                (2, atmosphere_transforms.uniforms().binding().unwrap()),
+                (3, view_uniforms.uniforms.binding().unwrap()),
+                (4, lights_uniforms.view_gpu_lights.binding().unwrap()),
+                (8, &textures.transmittance_lut.default_view),
+                (9, &samplers.transmittance_lut),
+                (10, &textures.multiscattering_lut.default_view),
+                (11, &samplers.multiscattering_lut),
+                (12, &textures.sky_view_lut.default_view),
+                (13, &samplers.sky_view_lut),
+                (14, &textures.aerial_view_lut.default_view),
+                (15, &samplers.aerial_view_lut),
+                (16, &textures.environment),
             )),
         );
 
@@ -135,7 +140,7 @@ pub(super) fn prepare_atmosphere_probe_bind_groups(
 }
 
 pub(super) fn prepare_probe_textures(
-    view_textures: Query<&AtmosphereTextures, With<Atmosphere>>,
+    view_textures: Query<&AtmosphereTextures, With<ExtractedAtmosphere>>,
     probes: Query<
         (Entity, &AtmosphereEnvironmentMap),
         (
@@ -246,7 +251,7 @@ pub fn prepare_atmosphere_probe_components(
 
 pub(super) struct EnvironmentNode {
     main_view_query: QueryState<(
-        Read<DynamicUniformIndex<Atmosphere>>,
+        Read<DynamicUniformIndex<GpuAtmosphere>>,
         Read<DynamicUniformIndex<GpuAtmosphereSettings>>,
         Read<AtmosphereTransformsOffset>,
         Read<ViewUniformOffset>,

crates/bevy_pbr/src/atmosphere/environment.wgsl

@@ -5,7 +5,7 @@
     utils::sample_cube_dir
 }
 
-@group(0) @binding(13) var output: texture_storage_2d_array<rgba16float, write>;
+@group(0) @binding(16) var output: texture_storage_2d_array<rgba16float, write>;
 
 @compute @workgroup_size(8, 8, 1)
 fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
@@ -36,4 +36,4 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
     let color = vec4<f32>(inscattering, 1.0);
 
     textureStore(output, vec2<i32>(global_id.xy), i32(slice_index), color);
-}
+}