Skip to main content

bevy_pbr/
medium.rs

1use bevy_color::ColorToComponents;
2use bevy_light::atmosphere::{ScatteringMedium, ScatteringTerm};
3
4use bevy_app::{App, Plugin};
5use bevy_asset::AssetId;
6use bevy_ecs::{
7    resource::Resource,
8    system::{Commands, Res, SystemParamItem},
9};
10use bevy_math::{ops, Vec4};
11use bevy_render::{
12    render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin},
13    render_resource::{
14        Extent3d, FilterMode, Sampler, SamplerDescriptor, Texture, TextureDataOrder,
15        TextureDescriptor, TextureDimension, TextureFormat, TextureUsages, TextureView,
16        TextureViewDescriptor,
17    },
18    renderer::{RenderDevice, RenderQueue},
19    RenderApp, RenderStartup,
20};
21use smallvec::SmallVec;
22
23#[doc(hidden)]
24pub struct ScatteringMediumPlugin;
25
26impl Plugin for ScatteringMediumPlugin {
27    fn build(&self, app: &mut App) {
28        app.add_plugins(RenderAssetPlugin::<GpuScatteringMedium>::default());
29
30        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
31            render_app.add_systems(RenderStartup, init_scattering_medium_sampler);
32        }
33    }
34}
35
36/// The GPU representation of a [`ScatteringMedium`].
37pub struct GpuScatteringMedium {
38    /// The terms of the scattering medium.
39    pub terms: SmallVec<[ScatteringTerm; 1]>,
40    /// The resolution at which to sample the falloff distribution of each
41    /// scattering term.
42    pub falloff_resolution: u32,
43    /// The resolution at which to sample the phase function of each
44    /// scattering term.
45    pub phase_resolution: u32,
46    /// The `density_lut`, a 2D `falloff_resolution x 2` LUT which contains the
47    /// medium's optical density with respect to the atmosphere's "falloff parameter",
48    /// a linear value which is 1.0 at the planet's surface and 0.0 at the edge of
49    /// space. The first and second rows correspond to absorption density and
50    /// scattering density respectively.
51    pub density_lut: Texture,
52    /// The default [`TextureView`] of the `density_lut`
53    pub density_lut_view: TextureView,
54    /// The `scattering_lut`, a 2D `falloff_resolution x phase_resolution` LUT which
55    /// contains the medium's scattering density multiplied by the phase function, with
56    /// the U axis corresponding to the falloff parameter and the V axis corresponding
57    /// to a nonlinear mapping of `neg_LdotV`, where `neg_LdotV` is the dot product of
58    /// the light direction and the incoming view vector.
59    pub scattering_lut: Texture,
60    /// The default [`TextureView`] of the `scattering_lut`
61    pub scattering_lut_view: TextureView,
62}
63
64impl RenderAsset for GpuScatteringMedium {
65    type SourceAsset = ScatteringMedium;
66
67    type Param = (Res<'static, RenderDevice>, Res<'static, RenderQueue>);
68
69    fn prepare_asset(
70        source_asset: Self::SourceAsset,
71        _asset_id: AssetId<Self::SourceAsset>,
72        (render_device, render_queue): &mut SystemParamItem<Self::Param>,
73        _previous_asset: Option<&Self>,
74    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
75        let mut density: Vec<Vec4> =
76            Vec::with_capacity(2 * source_asset.falloff_resolution as usize);
77
78        density.extend((0..source_asset.falloff_resolution).map(|i| {
79            let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
80
81            source_asset
82                .terms
83                .iter()
84                .map(|term| term.absorption.extend(0.0) * term.falloff.sample(falloff))
85                .sum::<Vec4>()
86        }));
87
88        density.extend((0..source_asset.falloff_resolution).map(|i| {
89            let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
90
91            source_asset
92                .terms
93                .iter()
94                .map(|term| term.scattering.extend(0.0) * term.falloff.sample(falloff))
95                .sum::<Vec4>()
96        }));
97
98        let mut scattering: Vec<Vec4> = Vec::with_capacity(
99            source_asset.falloff_resolution as usize * source_asset.phase_resolution as usize,
100        );
101
102        // Define the power curve parameter for the nonlinear phase mapping
103        const PHASE_MAPPING_N: f32 = 0.5;
104        let inv_n = 1.0 / PHASE_MAPPING_N;
105
106        scattering.extend(
107            (0..source_asset.falloff_resolution * source_asset.phase_resolution).map(|raw_i| {
108                let i = raw_i % source_asset.phase_resolution;
109                let j = raw_i / source_asset.phase_resolution;
110                let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
111                let phase = (j as f32 + 0.5) / source_asset.phase_resolution as f32;
112
113                // Nonlinear phase mapping to mitigate banding in low-resolution LUTs.
114                // Phase functions peak at forward/backward scattering. Linearly
115                // sampling these regions causes banding. The power curve (n < 1)
116                // allocates more texels near the peaks.
117                let s = 2.0 * phase - 1.0;
118                let neg_l_dot_v = s.signum() * (1.0 - ops::powf(1.0 - s.abs(), inv_n));
119
120                source_asset
121                    .terms
122                    .iter()
123                    .filter_map(|term| {
124                        let f = term.falloff.sample(falloff);
125                        term.phase
126                            .sample(neg_l_dot_v)
127                            .map(|phase| (term.scattering * phase.to_vec3() * f).extend(0.0))
128                    })
129                    .sum::<Vec4>()
130            }),
131        );
132
133        let density_lut = render_device.create_texture_with_data(
134            render_queue,
135            &TextureDescriptor {
136                label: source_asset
137                    .label
138                    .as_deref()
139                    .map(|label| format!("{}_density_lut", label))
140                    .as_deref()
141                    .or(Some("scattering_medium_density_lut")),
142                size: Extent3d {
143                    width: source_asset.falloff_resolution,
144                    height: 2,
145                    depth_or_array_layers: 1,
146                },
147                mip_level_count: 1,
148                sample_count: 1,
149                dimension: TextureDimension::D2,
150                format: TextureFormat::Rgba32Float,
151                usage: TextureUsages::TEXTURE_BINDING,
152                view_formats: &[],
153            },
154            TextureDataOrder::LayerMajor,
155            bytemuck::cast_slice(density.as_slice()),
156        );
157
158        let density_lut_view = density_lut.create_view(&TextureViewDescriptor {
159            label: source_asset
160                .label
161                .as_deref()
162                .map(|label| format!("{}_density_lut_view", label))
163                .as_deref()
164                .or(Some("scattering_medium_density_lut_view")),
165            ..Default::default()
166        });
167
168        let scattering_lut = render_device.create_texture_with_data(
169            render_queue,
170            &TextureDescriptor {
171                label: source_asset
172                    .label
173                    .as_deref()
174                    .map(|label| format!("{}_scattering_lut", label))
175                    .as_deref()
176                    .or(Some("scattering_medium_scattering_lut")),
177                size: Extent3d {
178                    width: source_asset.falloff_resolution,
179                    height: source_asset.phase_resolution,
180                    depth_or_array_layers: 1,
181                },
182                mip_level_count: 1,
183                sample_count: 1,
184                dimension: TextureDimension::D2,
185                format: TextureFormat::Rgba32Float,
186                usage: TextureUsages::TEXTURE_BINDING,
187                view_formats: &[],
188            },
189            TextureDataOrder::LayerMajor,
190            bytemuck::cast_slice(scattering.as_slice()),
191        );
192
193        let scattering_lut_view = scattering_lut.create_view(&TextureViewDescriptor {
194            label: source_asset
195                .label
196                .as_deref()
197                .map(|label| format!("{}_scattering_lut", label))
198                .as_deref()
199                .or(Some("scattering_medium_scattering_lut_view")),
200            ..Default::default()
201        });
202
203        Ok(Self {
204            terms: source_asset.terms,
205            falloff_resolution: source_asset.falloff_resolution,
206            phase_resolution: source_asset.phase_resolution,
207            density_lut,
208            density_lut_view,
209            scattering_lut,
210            scattering_lut_view,
211        })
212    }
213}
214
215/// The default sampler for all scattering media LUTs.
216///
217/// Just a bilinear clamp-to-edge sampler, nothing fancy.
218#[derive(Resource)]
219pub struct ScatteringMediumSampler(Sampler);
220
221impl ScatteringMediumSampler {
222    pub fn sampler(&self) -> &Sampler {
223        &self.0
224    }
225}
226
227fn init_scattering_medium_sampler(mut commands: Commands, render_device: Res<RenderDevice>) {
228    let sampler = render_device.create_sampler(&SamplerDescriptor {
229        label: Some("scattering_medium_sampler"),
230        mag_filter: FilterMode::Linear,
231        min_filter: FilterMode::Linear,
232        ..Default::default()
233    });
234
235    commands.insert_resource(ScatteringMediumSampler(sampler));
236}