bevy_pbr/
medium.rs

1use alloc::{borrow::Cow, sync::Arc};
2use core::f32::{self, consts::PI};
3
4use bevy_app::{App, Plugin};
5use bevy_asset::{Asset, AssetApp, AssetId};
6use bevy_ecs::{
7    resource::Resource,
8    system::{Commands, Res, SystemParamItem},
9};
10use bevy_math::{ops, Curve, FloatPow, Vec3, Vec4};
11use bevy_reflect::TypePath;
12use bevy_render::{
13    render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin},
14    render_resource::{
15        Extent3d, FilterMode, Sampler, SamplerDescriptor, Texture, TextureDataOrder,
16        TextureDescriptor, TextureDimension, TextureFormat, TextureUsages, TextureView,
17        TextureViewDescriptor,
18    },
19    renderer::{RenderDevice, RenderQueue},
20    RenderApp, RenderStartup,
21};
22use smallvec::SmallVec;
23
24#[doc(hidden)]
25pub struct ScatteringMediumPlugin;
26
27impl Plugin for ScatteringMediumPlugin {
28    fn build(&self, app: &mut App) {
29        app.init_asset::<ScatteringMedium>()
30            .add_plugins(RenderAssetPlugin::<GpuScatteringMedium>::default());
31
32        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
33            render_app.add_systems(RenderStartup, init_scattering_medium_sampler);
34        }
35    }
36}
37
38/// An asset that defines how a material scatters light.
39///
40/// In order to calculate how light passes through a medium,
41/// you need three pieces of information:
42/// - how much light the medium *absorbs* per unit length
43/// - how much light the medium *scatters* per unit length
44/// - what *directions* the medium is likely to scatter light in.
45///
46/// The first two are fairly simple, and are sometimes referred to together
47/// (accurately enough) as the medium's [optical density].
48///
49/// The last, defined by a [phase function], is the most important in creating
50/// the look of a medium. Our brains are very good at noticing (if unconsciously)
51/// that a dust storm scatters light differently than a rain cloud, for example.
52/// See the docs on [`PhaseFunction`] for more info.
53///
54/// In reality, media are often composed of multiple elements that scatter light
55/// independently, for Earth's atmosphere is composed of the gas itself, but also
56/// suspended dust and particulate. These each scatter light differently, and are
57/// distributed in different amounts at different altitudes. In a [`ScatteringMedium`],
58/// these are each represented by a [`ScatteringTerm`]
59///
60/// ## Technical Details
61///
62/// A [`ScatteringMedium`] is represented on the GPU by a set of two LUTs, which
63/// are re-created every time the asset is modified. See the docs on
64/// [`GpuScatteringMedium`] for more info.
65///
66/// [optical density]: https://en.wikipedia.org/wiki/Optical_Density
67/// [phase function]: https://www.pbr-book.org/4ed/Volume_Scattering/Phase_Functions
68#[derive(TypePath, Asset, Clone)]
69pub struct ScatteringMedium {
70    /// An optional label for the medium, used when creating the LUTs on the GPU.
71    pub label: Option<Cow<'static, str>>,
72    /// The resolution at which to sample the falloff distribution of each
73    /// scattering term. Custom or more detailed distributions may benefit
74    /// from a higher value, at the cost of more memory use.
75    pub falloff_resolution: u32,
76    /// The resolution at which to sample the phase function of each scattering
77    /// term. Custom or more detailed phase functions may benefit from a higher
78    /// value, at the cost of more memory use.
79    pub phase_resolution: u32,
80    /// The list of [`ScatteringTerm`]s that compose this [`ScatteringMedium`]
81    pub terms: SmallVec<[ScatteringTerm; 1]>,
82}
83
84impl Default for ScatteringMedium {
85    fn default() -> Self {
86        ScatteringMedium::earthlike(256, 256)
87    }
88}
89
90impl ScatteringMedium {
91    // Returns a scattering medium with a default label and the
92    // specified scattering terms.
93    pub fn new(
94        falloff_resolution: u32,
95        phase_resolution: u32,
96        terms: impl IntoIterator<Item = ScatteringTerm>,
97    ) -> Self {
98        Self {
99            label: None,
100            falloff_resolution,
101            phase_resolution,
102            terms: terms.into_iter().collect(),
103        }
104    }
105
106    // Consumes and returns this scattering medium with a new label.
107    pub fn with_label(self, label: impl Into<Cow<'static, str>>) -> Self {
108        Self {
109            label: Some(label.into()),
110            ..self
111        }
112    }
113
114    // Consumes and returns this scattering medium with each scattering terms'
115    // densities multiplied by `multiplier`.
116    pub fn with_density_multiplier(mut self, multiplier: f32) -> Self {
117        self.terms.iter_mut().for_each(|term| {
118            term.absorption *= multiplier;
119            term.scattering *= multiplier;
120        });
121
122        self
123    }
124
125    /// Returns a scattering medium representing an earthlike atmosphere.
126    ///
127    /// Uses physically-based scale heights from Earth's atmosphere, assuming
128    /// a 60 km atmosphere height:
129    /// - Rayleigh (molecular) scattering: 8 km scale height
130    /// - Mie (aerosol) scattering: 1.2 km scale height
131    pub fn earthlike(falloff_resolution: u32, phase_resolution: u32) -> Self {
132        Self::new(
133            falloff_resolution,
134            phase_resolution,
135            [
136                // Rayleigh scattering Term
137                ScatteringTerm {
138                    absorption: Vec3::ZERO,
139                    scattering: Vec3::new(5.802e-6, 13.558e-6, 33.100e-6),
140                    falloff: Falloff::Exponential { scale: 8.0 / 60.0 },
141                    phase: PhaseFunction::Rayleigh,
142                },
143                // Mie scattering Term
144                ScatteringTerm {
145                    absorption: Vec3::splat(3.996e-6),
146                    scattering: Vec3::splat(0.444e-6),
147                    falloff: Falloff::Exponential { scale: 1.2 / 60.0 },
148                    phase: PhaseFunction::Mie { asymmetry: 0.8 },
149                },
150                // Ozone scattering Term
151                ScatteringTerm {
152                    absorption: Vec3::new(0.650e-6, 1.881e-6, 0.085e-6),
153                    scattering: Vec3::ZERO,
154                    falloff: Falloff::Tent {
155                        center: 0.75,
156                        width: 0.3,
157                    },
158                    phase: PhaseFunction::Isotropic,
159                },
160            ],
161        )
162        .with_label("earthlike_atmosphere")
163    }
164}
165
166/// An individual element of a [`ScatteringMedium`].
167///
168/// A [`ScatteringMedium`] can be built out of a number of simpler [`ScatteringTerm`]s,
169/// which correspond to an individual element of the medium. For example, Earth's
170/// atmosphere would be (roughly) composed of two [`ScatteringTerm`]s: the atmospheric
171/// gases themselves, which extend to the edge of space, and suspended dust particles,
172/// which are denser but lie closer to the ground.
173#[derive(Default, Clone)]
174pub struct ScatteringTerm {
175    /// This term's optical absorption density, or how much light of each wavelength
176    /// it absorbs per meter.
177    ///
178    /// units: m^-1
179    pub absorption: Vec3,
180    /// This term's optical scattering density, or how much light of each wavelength
181    /// it scatters per meter.
182    ///
183    /// units: m^-1
184    pub scattering: Vec3,
185    /// This term's falloff distribution. See the docs on [`Falloff`] for more info.
186    pub falloff: Falloff,
187    /// This term's [phase function], which determines the character of how it
188    /// scatters light. See the docs on [`PhaseFunction`] for more info.
189    ///
190    /// [phase function]: https://www.pbr-book.org/4ed/Volume_Scattering/Phase_Functions
191    pub phase: PhaseFunction,
192}
193
194/// Describes how the media in a [`ScatteringTerm`] is distributed.
195///
196/// This is closely related to the optical density values [`ScatteringTerm::absorption`] and
197/// [`ScatteringTerm::scattering`]. Most media aren't the same density everywhere;
198/// near the edge of space Earth's atmosphere is much less dense, and it absorbs
199/// and scatters less light.
200///
201/// [`Falloff`] determines how the density of a medium changes as a function of
202/// an abstract "falloff parameter" `p`. `p = 1` denotes where the medium is the
203/// densest, i.e. at the surface of the Earth, `p = 0` denotes where the medium
204/// fades away completely, i.e. at the edge of space, and values between scale
205/// linearly with distance, so `p = 0.5` would be halfway between the surface
206/// and the edge of space.
207///
208/// When processing a [`ScatteringMedium`], the `absorption` and `scattering` values
209/// for each [`ScatteringTerm`] are multiplied by the value of the falloff function, `f(p)`.
210#[derive(Default, Clone)]
211pub enum Falloff {
212    /// A simple linear falloff function, which essentially
213    /// passes the falloff parameter through unchanged.
214    ///
215    /// f(1) = 1
216    /// f(0) = 0
217    /// f(p) = p
218    #[default]
219    Linear,
220    /// An exponential falloff function parametrized by a proportional scale.
221    /// When paired with an absolute "falloff distance" like the distance from
222    /// Earth's surface to the edge of space, this is analogous to the "height
223    /// scale" value common in atmospheric scattering literature, though it will
224    /// diverge from this for large or negative `scale` values.
225    ///
226    /// f(1) = 1
227    /// f(0) = 0
228    /// f(p) = (e^((1-p)/s) - e^(1/s))/(e - e^(1/s))
229    Exponential {
230        /// The "scale" of the exponential falloff. Values closer to zero will
231        /// produce steeper falloff, and values farther from zero will produce
232        /// gentler falloff, approaching linear falloff as scale goes to `+-∞`.
233        ///
234        /// Negative values change the *concavity* of the falloff function:
235        /// rather than an initial narrow region of steep falloff followed by a
236        /// wide region of gentle falloff, there will be an initial wide region
237        /// of gentle falloff followed by a narrow region of steep falloff.
238        ///
239        /// domain: (-∞, ∞)
240        ///
241        /// NOTE, this function is not defined when `scale == 0`.
242        /// In that case, it will fall back to linear falloff.
243        scale: f32,
244    },
245    /// A tent-shaped falloff function, which produces a triangular
246    /// peak at the center and linearly falls off to either side.
247    ///
248    /// f(`center`) = 1
249    /// f(`center` +- `width` / 2) = 0
250    Tent {
251        /// The center of the tent function peak
252        ///
253        /// domain: [0, 1]
254        center: f32,
255        /// The total width of the tent function peak
256        ///
257        /// domain: [0, 1]
258        width: f32,
259    },
260    /// A falloff function defined by a custom curve.
261    ///
262    /// domain: [0, 1],
263    /// range: [0, 1],
264    Curve(Arc<dyn Curve<f32> + Send + Sync>),
265}
266
267impl Falloff {
268    /// Returns a falloff function corresponding to a custom curve.
269    pub fn from_curve(curve: impl Curve<f32> + Send + Sync + 'static) -> Self {
270        Self::Curve(Arc::new(curve))
271    }
272
273    fn sample(&self, p: f32) -> f32 {
274        match self {
275            Falloff::Linear => p,
276            Falloff::Exponential { scale } => {
277                // fill discontinuity at scale == 0,
278                // arbitrarily choose linear falloff
279                if *scale == 0.0 {
280                    p
281                } else {
282                    let s = -1.0 / scale;
283                    let exp_p_s = ops::exp((1.0 - p) * s);
284                    let exp_s = ops::exp(s);
285                    (exp_p_s - exp_s) / (1.0 - exp_s)
286                }
287            }
288            Falloff::Tent { center, width } => (1.0 - (p - center).abs() / (0.5 * width)).max(0.0),
289            Falloff::Curve(curve) => curve.sample(p).unwrap_or(0.0),
290        }
291    }
292}
293
294/// Describes how a [`ScatteringTerm`] scatters light in different directions.
295///
296/// A [phase function] is a function `f: [-1, 1] -> [0, ∞)`, symmetric about `x=0`
297/// whose input is the cosine of the angle between an incoming light direction and
298/// and outgoing light direction, and whose output is the proportion of the incoming
299/// light that is actually scattered in that direction.
300///
301/// The phase function has an important effect on the "look" of a medium in a scene.
302/// Media consisting of particles of a different size or shape scatter light differently,
303/// and our brains are very good at telling the difference. A dust cloud, which might
304/// correspond roughly to `PhaseFunction::Mie { asymmetry: 0.8 }`, looks quite different
305/// from the rest of the sky (atmospheric gases), which correspond to `PhaseFunction::Rayleigh`
306///
307/// [phase function]: https://www.pbr-book.org/4ed/Volume_Scattering/Phase_Functions
308#[derive(Clone)]
309pub enum PhaseFunction {
310    /// A phase function that scatters light evenly in all directions.
311    Isotropic,
312
313    /// A phase function representing [Rayleigh scattering].
314    ///
315    /// Rayleigh scattering occurs naturally for particles much smaller than
316    /// the wavelengths of visible light, such as gas molecules in the atmosphere.
317    /// It's generally wavelength-dependent, where shorter wavelengths are scattered
318    /// more strongly, so [scattering](ScatteringTerm::scattering) should have
319    /// higher values for blue than green and green than red. Particles that
320    /// participate in Rayleigh scattering don't absorb any light, either.
321    ///
322    /// [Rayleigh scattering]: https://en.wikipedia.org/wiki/Rayleigh_scattering
323    Rayleigh,
324
325    /// The [Henyey-Greenstein phase function], which approximates [Mie scattering].
326    ///
327    /// Mie scattering occurs naturally for spherical particles of dust
328    /// and aerosols roughly the same size as the wavelengths of visible light,
329    /// so it's useful for representing dust or sea spray. It's generally
330    /// wavelength-independent, so [absorption](ScatteringTerm::absorption)
331    /// and [scattering](ScatteringTerm::scattering) should be set to a greyscale value.
332    ///
333    /// [Mie scattering]: https://en.wikipedia.org/wiki/Mie_scattering
334    /// [Henyey-Greenstein phase function]: https://www.oceanopticsbook.info/view/scattering/level-2/the-henyey-greenstein-phase-function
335    Mie {
336        /// Whether the Mie scattering function is biased towards scattering
337        /// light forwards (asymmetry > 0) or backwards (asymmetry < 0).
338        ///
339        /// domain: [-1, 1]
340        asymmetry: f32,
341    },
342
343    /// A phase function defined by a custom curve, where the input
344    /// is the cosine of the angle between the incoming light ray
345    /// and the scattered light ray, and the output is the fraction
346    /// of the incoming light scattered in that direction.
347    ///
348    /// Note: it's important for photorealism that the phase function
349    /// be *energy conserving*, meaning that in total no more light can
350    /// be scattered than actually entered the medium. For this to be
351    /// the case, the integral of the phase function over its domain must
352    /// be equal to 1/2π.
353    ///
354    ///   1
355    /// ∫   p(x) dx = 1/2π
356    ///  -1
357    ///
358    /// domain: [-1, 1]
359    /// range: [0, 1]
360    Curve(Arc<dyn Curve<f32> + Send + Sync>),
361}
362
363impl PhaseFunction {
364    /// A phase function defined by a custom curve.
365    pub fn from_curve(curve: impl Curve<f32> + Send + Sync + 'static) -> Self {
366        Self::Curve(Arc::new(curve))
367    }
368
369    fn sample(&self, neg_l_dot_v: f32) -> f32 {
370        const FRAC_4_PI: f32 = 0.25 / PI;
371        const FRAC_3_16_PI: f32 = 0.1875 / PI;
372        match self {
373            PhaseFunction::Isotropic => FRAC_4_PI,
374            PhaseFunction::Rayleigh => FRAC_3_16_PI * (1.0 + neg_l_dot_v * neg_l_dot_v),
375            PhaseFunction::Mie { asymmetry } => {
376                let denom = 1.0 + asymmetry.squared() - 2.0 * asymmetry * neg_l_dot_v;
377                FRAC_4_PI * (1.0 - asymmetry.squared()) / (denom * denom.sqrt())
378            }
379            PhaseFunction::Curve(curve) => curve.sample(neg_l_dot_v).unwrap_or(0.0),
380        }
381    }
382}
383
384impl Default for PhaseFunction {
385    fn default() -> Self {
386        Self::Mie { asymmetry: 0.8 }
387    }
388}
389
390/// The GPU representation of a [`ScatteringMedium`].
391pub struct GpuScatteringMedium {
392    /// The terms of the scattering medium.
393    pub terms: SmallVec<[ScatteringTerm; 1]>,
394    /// The resolution at which to sample the falloff distribution of each
395    /// scattering term.
396    pub falloff_resolution: u32,
397    /// The resolution at which to sample the phase function of each
398    /// scattering term.
399    pub phase_resolution: u32,
400    /// The `density_lut`, a 2D `falloff_resolution x 2` LUT which contains the
401    /// medium's optical density with respect to the atmosphere's "falloff parameter",
402    /// a linear value which is 1.0 at the planet's surface and 0.0 at the edge of
403    /// space. The first and second rows correspond to absorption density and
404    /// scattering density respectively.
405    pub density_lut: Texture,
406    /// The default [`TextureView`] of the `density_lut`
407    pub density_lut_view: TextureView,
408    /// The `scattering_lut`, a 2D `falloff_resolution x phase_resolution` LUT which
409    /// contains the medium's scattering density multiplied by the phase function, with
410    /// the U axis corresponding to the falloff parameter and the V axis corresponding
411    /// to `neg_LdotV * 0.5 + 0.5`, where `neg_LdotV` is the dot product of the light
412    /// direction and the incoming view vector.
413    pub scattering_lut: Texture,
414    /// The default [`TextureView`] of the `scattering_lut`
415    pub scattering_lut_view: TextureView,
416}
417
418impl RenderAsset for GpuScatteringMedium {
419    type SourceAsset = ScatteringMedium;
420
421    type Param = (Res<'static, RenderDevice>, Res<'static, RenderQueue>);
422
423    fn prepare_asset(
424        source_asset: Self::SourceAsset,
425        _asset_id: AssetId<Self::SourceAsset>,
426        (render_device, render_queue): &mut SystemParamItem<Self::Param>,
427        _previous_asset: Option<&Self>,
428    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
429        let mut density: Vec<Vec4> =
430            Vec::with_capacity(2 * source_asset.falloff_resolution as usize);
431
432        density.extend((0..source_asset.falloff_resolution).map(|i| {
433            let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
434
435            source_asset
436                .terms
437                .iter()
438                .map(|term| term.absorption.extend(0.0) * term.falloff.sample(falloff))
439                .sum::<Vec4>()
440        }));
441
442        density.extend((0..source_asset.falloff_resolution).map(|i| {
443            let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
444
445            source_asset
446                .terms
447                .iter()
448                .map(|term| term.scattering.extend(0.0) * term.falloff.sample(falloff))
449                .sum::<Vec4>()
450        }));
451
452        let mut scattering: Vec<Vec4> = Vec::with_capacity(
453            source_asset.falloff_resolution as usize * source_asset.phase_resolution as usize,
454        );
455
456        scattering.extend(
457            (0..source_asset.falloff_resolution * source_asset.phase_resolution).map(|raw_i| {
458                let i = raw_i % source_asset.phase_resolution;
459                let j = raw_i / source_asset.phase_resolution;
460                let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
461                let phase = (j as f32 + 0.5) / source_asset.phase_resolution as f32;
462                let neg_l_dot_v = phase * 2.0 - 1.0;
463
464                source_asset
465                    .terms
466                    .iter()
467                    .map(|term| {
468                        term.scattering.extend(0.0)
469                            * term.falloff.sample(falloff)
470                            * term.phase.sample(neg_l_dot_v)
471                    })
472                    .sum::<Vec4>()
473            }),
474        );
475
476        let density_lut = render_device.create_texture_with_data(
477            render_queue,
478            &TextureDescriptor {
479                label: source_asset
480                    .label
481                    .as_deref()
482                    .map(|label| format!("{}_density_lut", label))
483                    .as_deref()
484                    .or(Some("scattering_medium_density_lut")),
485                size: Extent3d {
486                    width: source_asset.falloff_resolution,
487                    height: 2,
488                    depth_or_array_layers: 1,
489                },
490                mip_level_count: 1,
491                sample_count: 1,
492                dimension: TextureDimension::D2,
493                format: TextureFormat::Rgba32Float,
494                usage: TextureUsages::TEXTURE_BINDING,
495                view_formats: &[],
496            },
497            TextureDataOrder::LayerMajor,
498            bytemuck::cast_slice(density.as_slice()),
499        );
500
501        let density_lut_view = density_lut.create_view(&TextureViewDescriptor {
502            label: source_asset
503                .label
504                .as_deref()
505                .map(|label| format!("{}_density_lut_view", label))
506                .as_deref()
507                .or(Some("scattering_medium_density_lut_view")),
508            ..Default::default()
509        });
510
511        let scattering_lut = render_device.create_texture_with_data(
512            render_queue,
513            &TextureDescriptor {
514                label: source_asset
515                    .label
516                    .as_deref()
517                    .map(|label| format!("{}_scattering_lut", label))
518                    .as_deref()
519                    .or(Some("scattering_medium_scattering_lut")),
520                size: Extent3d {
521                    width: source_asset.falloff_resolution,
522                    height: source_asset.phase_resolution,
523                    depth_or_array_layers: 1,
524                },
525                mip_level_count: 1,
526                sample_count: 1,
527                dimension: TextureDimension::D2,
528                format: TextureFormat::Rgba32Float,
529                usage: TextureUsages::TEXTURE_BINDING,
530                view_formats: &[],
531            },
532            TextureDataOrder::LayerMajor,
533            bytemuck::cast_slice(scattering.as_slice()),
534        );
535
536        let scattering_lut_view = scattering_lut.create_view(&TextureViewDescriptor {
537            label: source_asset
538                .label
539                .as_deref()
540                .map(|label| format!("{}_scattering_lut", label))
541                .as_deref()
542                .or(Some("scattering_medium_scattering_lut_view")),
543            ..Default::default()
544        });
545
546        Ok(Self {
547            terms: source_asset.terms,
548            falloff_resolution: source_asset.falloff_resolution,
549            phase_resolution: source_asset.phase_resolution,
550            density_lut,
551            density_lut_view,
552            scattering_lut,
553            scattering_lut_view,
554        })
555    }
556}
557
558/// The default sampler for all scattering media LUTs.
559///
560/// Just a bilinear clamp-to-edge sampler, nothing fancy.
561#[derive(Resource)]
562pub struct ScatteringMediumSampler(Sampler);
563
564impl ScatteringMediumSampler {
565    pub fn sampler(&self) -> &Sampler {
566        &self.0
567    }
568}
569
570fn init_scattering_medium_sampler(mut commands: Commands, render_device: Res<RenderDevice>) {
571    let sampler = render_device.create_sampler(&SamplerDescriptor {
572        label: Some("scattering_medium_sampler"),
573        mag_filter: FilterMode::Linear,
574        min_filter: FilterMode::Linear,
575        ..Default::default()
576    });
577
578    commands.insert_resource(ScatteringMediumSampler(sampler));
579}