Generalized atmospheric scattering media (#20838)

# Objective

Right now, only a very small set of atmospheres are possible with Bevy's
atmospheric scattering system because of the fixed set of scattering
terms available. For example, a scene set in a dry, desert environment
might want a dense low-lying layer of dust particulate separate from the
ambient dust in the atmosphere. This PR introduces a mechanism for
generalized scattering media, replacing the fixed scattering terms with
a customizable asset.

## Solution

```rust
#[derive(TypePath, Asset, Clone)]
pub struct ScatteringMedium {
    pub label: Option<Cow<'static, str>>,
    pub falloff_resolution: u32,
    pub phase_resolution: u32,
    pub terms: SmallVec<[ScatteringTerm; 1]>,
}
```

<details>
  <summary>see other new types</summary>

```rust
#[derive(Default, Clone)]
pub struct ScatteringTerm {
    pub absorption: Vec3,
    pub scattering: Vec3,
    pub falloff: Falloff,
    pub phase: PhaseFunction,
}

#[derive(Default, Clone)]
pub enum Falloff {
    #[default]
    Linear,
    Exponential {
        scale: f32,
    },
    Tent {
        center: f32,
        width: f32,
    },
    /// A falloff function defined by a custom curve.
    ///
    /// domain: [0, 1),
    /// range: [0, 1],
    Curve(Arc<dyn Curve<f32> + Send + Sync>),
}

#[derive(Clone)]
pub enum PhaseFunction {
    Isotropic,
    Rayleigh,
    Mie {
        /// domain: [-1, 1]
        bias: f32,
    },
    /// A phase function defined by a custom curve.
    ///
    /// domain: [-1, 1]
    /// range: [0, 1]
    Curve(Arc<dyn Curve<f32> + Send + Sync>),
}
```
</details>

`ScatteringMedium` contains a list of `ScatteringTerms`, which are
processed into a set of two LUTs:

- The "density LUT", a 2D `falloff_resolution x 2` LUT which contains
the medium's optical density with respect to the atmosphere's "falloff
parameter", a linear value which is 1.0 at the planet's surface and 0.0
at the edge of space. Absorption density and scattering density
correspond to the first and second rows respectively.
- The "scattering LUT", a 2D `falloff_resolution x phase_resolution` LUT
which contains the medium's scattering density multiplied by the phase
function, with the U axis corresponding to the falloff parameter and the
V axis corresponding to `neg_LdotV * 0.5 + 0.5`, where `neg_LdotV` is
the dot product of the light direction and the outgoing view vector.

## Testing

- Need to verify output, should be almost exactly the same
- exponential falloff is slightly different now, but verified new
parameters against the old in Desmos.

## TODOS: 

- Docs Docs Docs
- Cleanup
- profile perf
- reduce memory usage/traffic. This approach requires a few extra
texture samples in the inner loop of each pass. Each atmosphere LUT is
still quite small, but texture samples are expensive and the new LUTs
use f32 texels currently.

## Showcase

<details>
  <summary>Click to view showcase</summary>

```rust
fn init_atmosphere(mut commands: Commands, scattering_media: ResMut<Assets<ScatteringMedium>>) {
  let earth_atmosphere = scattering_media.add(
    ScatteringMedium::new(
      256,
      256,
      [
          // rayleigh scattering
          ScatteringTerm {
              absorption: Vec3::ZERO,
              scattering: Vec3::new(5.802e-6, 13.558e-6, 33.100e-6),
              falloff: Falloff::Exponential { scale: 12.5 },
              phase: PhaseFunction::Rayleigh,
          },
          // mie scattering
          ScatteringTerm {
              absorption: Vec3::splat(3.996e-6),
              scattering: Vec3::splat(0.444e-6),
              falloff: Falloff::Exponential { scale: 83.5 },
              phase: PhaseFunction::Mie { bias: 0.8 },
          },
          // ozone
          ScatteringTerm {
              absorption: Vec3::new(0.650e-6, 1.881e-6, 0.085e-6),
              scattering: Vec3::ZERO,
              falloff: Falloff::Tent {
                  center: 0.75,
                  width: 0.3,
              },
              phase: PhaseFunction::Isotropic,
          },
      ],
  ));

  commands.spawn((
    Camera3d::default(), 
    Atmosphere {
      bottom_radius: 6_360_000.0,
      top_radius: 6_460_000.0,
      ground_albedo: Vec3::splat(0.3),
      medium: earth_atmosphere,
    },
  ));
}
```
</details>

---------

Co-authored-by: Máté Homolya <[email protected]>

Author: ecoskey

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);
-}
+}