1#define_import_path bevy_pbr::atmosphere::functions
2
3#import bevy_render::maths::{PI, HALF_PI, PI_2, fast_acos, fast_acos_4, fast_atan2}
4
5#import bevy_pbr::atmosphere::{
6    types::Atmosphere,
7    bindings::{
8        atmosphere, settings, view, lights, transmittance_lut, transmittance_lut_sampler, 
9        multiscattering_lut, multiscattering_lut_sampler, sky_view_lut, sky_view_lut_sampler,
10        aerial_view_lut, aerial_view_lut_sampler, atmosphere_transforms
11    },
12    bruneton_functions::{
13        transmittance_lut_r_mu_to_uv, transmittance_lut_uv_to_r_mu, 
14        ray_intersects_ground, distance_to_top_atmosphere_boundary, 
15        distance_to_bottom_atmosphere_boundary
16    },
17}
18
19// NOTE FOR CONVENTIONS: 
20// r:
21//   radius, or distance from planet center 
22//
23// altitude:
24//   distance from planet **surface**
25//
26// mu:
27//   cosine of the zenith angle of a ray with
28//   respect to the planet normal
29//
30// atmosphere space:
31//   abbreviated as "as" (contrast with vs, cs, ws), this space is similar
32//   to view space, but with the camera positioned horizontally on the planet
33//   surface, so the horizon is a horizontal line centered vertically in the
34//   frame. This enables the non-linear latitude parametrization the paper uses 
35//   to concentrate detail near the horizon 
36
37
38// CONSTANTS
39
40const FRAC_PI: f32 = 0.3183098862; // 1 / π
41const FRAC_2_PI: f32 = 0.15915494309;  // 1 / (2π)
42const FRAC_3_16_PI: f32 = 0.0596831036594607509; // 3 / (16π)
43const FRAC_4_PI: f32 = 0.07957747154594767; // 1 / (4π)
44const ROOT_2: f32 = 1.41421356; // √2
45
46// During raymarching, each segment is sampled at a single point. This constant determines
47// where in the segment that sample is taken (0.0 = start, 0.5 = middle, 1.0 = end).
48// We use 0.3 to sample closer to the start of each segment, which better approximates
49// the exponential falloff of atmospheric density.
50const MIDPOINT_RATIO: f32 = 0.3;
51
52// LUT UV PARAMATERIZATIONS
53
54fn unit_to_sub_uvs(val: vec2<f32>, resolution: vec2<f32>) -> vec2<f32> {
55    return (val + 0.5f / resolution) * (resolution / (resolution + 1.0f));
56}
57
58fn sub_uvs_to_unit(val: vec2<f32>, resolution: vec2<f32>) -> vec2<f32> {
59    return (val - 0.5f / resolution) * (resolution / (resolution - 1.0f));
60}
61
62fn multiscattering_lut_r_mu_to_uv(r: f32, mu: f32) -> vec2<f32> {
63    let u = 0.5 + 0.5 * mu;
64    let v = saturate((r - atmosphere.bottom_radius) / (atmosphere.top_radius - atmosphere.bottom_radius)); //TODO
65    return unit_to_sub_uvs(vec2(u, v), vec2<f32>(settings.multiscattering_lut_size));
66}
67
68fn multiscattering_lut_uv_to_r_mu(uv: vec2<f32>) -> vec2<f32> {
69    let adj_uv = sub_uvs_to_unit(uv, vec2<f32>(settings.multiscattering_lut_size));
70    let r = mix(atmosphere.bottom_radius, atmosphere.top_radius, adj_uv.y);
71    let mu = adj_uv.x * 2 - 1;
72    return vec2(r, mu);
73}
74
75fn sky_view_lut_r_mu_azimuth_to_uv(r: f32, mu: f32, azimuth: f32) -> vec2<f32> {
76    let u = (azimuth * FRAC_2_PI) + 0.5;
77
78    let v_horizon = sqrt(r * r - atmosphere.bottom_radius * atmosphere.bottom_radius);
79    let cos_beta = v_horizon / r;
80    // Using fast_acos_4 for better precision at small angles
81    // to avoid artifacts at the horizon
82    let beta = fast_acos_4(cos_beta);
83    let horizon_zenith = PI - beta;
84    let view_zenith = fast_acos_4(mu);
85
86    // Apply non-linear transformation to compress more texels 
87    // near the horizon where high-frequency details matter most
88    // l is latitude in [-π/2, π/2] and v is texture coordinate in [0,1]
89    let l = view_zenith - horizon_zenith;
90    let abs_l = abs(l);
91
92    let v = 0.5 + 0.5 * sign(l) * sqrt(abs_l / HALF_PI);
93
94    return unit_to_sub_uvs(vec2(u, v), vec2<f32>(settings.sky_view_lut_size));
95}
96
97fn sky_view_lut_uv_to_zenith_azimuth(r: f32, uv: vec2<f32>) -> vec2<f32> {
98    let adj_uv = sub_uvs_to_unit(vec2(uv.x, 1.0 - uv.y), vec2<f32>(settings.sky_view_lut_size));
99    let azimuth = (adj_uv.x - 0.5) * PI_2;
100
101    // Horizon parameters
102    let v_horizon = sqrt(r * r - atmosphere.bottom_radius * atmosphere.bottom_radius);
103    let cos_beta = v_horizon / r;
104    let beta = fast_acos_4(cos_beta);
105    let horizon_zenith = PI - beta;
106
107    // Inverse of horizon-detail mapping to recover original latitude from texture coordinate
108    let t = abs(2.0 * (adj_uv.y - 0.5));
109    let l = sign(adj_uv.y - 0.5) * HALF_PI * t * t;
110
111    return vec2(horizon_zenith - l, azimuth);
112}
113
114// LUT SAMPLING
115
116fn sample_transmittance_lut(r: f32, mu: f32) -> vec3<f32> {
117    let uv = transmittance_lut_r_mu_to_uv(r, mu);
118    return textureSampleLevel(transmittance_lut, transmittance_lut_sampler, uv, 0.0).rgb;
119}
120
121// NOTICE: This function is copyrighted by Eric Bruneton and INRIA, and falls
122// under the license reproduced in bruneton_functions.wgsl (variant of MIT license)
123//
124// FIXME: this function should be in bruneton_functions.wgsl, but because naga_oil doesn't 
125// support cyclic imports it's stuck here
126fn sample_transmittance_lut_segment(r: f32, mu: f32, t: f32) -> vec3<f32> {
127    let r_t = get_local_r(r, mu, t);
128    let mu_t = clamp((r * mu + t) / r_t, -1.0, 1.0);
129
130    if ray_intersects_ground(r, mu) {
131        return min(
132            sample_transmittance_lut(r_t, -mu_t) / sample_transmittance_lut(r, -mu),
133            vec3(1.0)
134        );
135    } else {
136        return min(
137            sample_transmittance_lut(r, mu) / sample_transmittance_lut(r_t, mu_t), vec3(1.0)
138        );
139    }
140}
141
142fn sample_multiscattering_lut(r: f32, mu: f32) -> vec3<f32> {
143    let uv = multiscattering_lut_r_mu_to_uv(r, mu);
144    return textureSampleLevel(multiscattering_lut, multiscattering_lut_sampler, uv, 0.0).rgb;
145}
146
147fn sample_sky_view_lut(r: f32, ray_dir_as: vec3<f32>) -> vec3<f32> {
148    let mu = ray_dir_as.y;
149    let azimuth = fast_atan2(ray_dir_as.x, -ray_dir_as.z);
150    let uv = sky_view_lut_r_mu_azimuth_to_uv(r, mu, azimuth);
151    return textureSampleLevel(sky_view_lut, sky_view_lut_sampler, uv, 0.0).rgb;
152}
153
154fn ndc_to_camera_dist(ndc: vec3<f32>) -> f32 {
155    let view_pos = view.view_from_clip * vec4(ndc, 1.0);
156    let t = length(view_pos.xyz / view_pos.w) * settings.scene_units_to_m;
157    return t;
158}
159
160// RGB channels: total inscattered light along the camera ray to the current sample.
161// A channel: average transmittance across all wavelengths to the current sample.
162fn sample_aerial_view_lut(uv: vec2<f32>, t: f32) -> vec3<f32> {
163    let t_max = settings.aerial_view_lut_max_distance;
164    let num_slices = f32(settings.aerial_view_lut_size.z);
165    // Each texel stores the value of the scattering integral over the whole slice,
166    // which requires us to offset the w coordinate by half a slice. For
167    // example, if we wanted the value of the integral at the boundary between slices,
168    // we'd need to sample at the center of the previous slice, and vice-versa for
169    // sampling in the center of a slice.
170    let uvw = vec3(uv, saturate(t / t_max - 0.5 / num_slices));
171    let sample = textureSampleLevel(aerial_view_lut, aerial_view_lut_sampler, uvw, 0.0);
172    // Since sampling anywhere between w=0 and w=t_slice will clamp to the first slice,
173    // we need to do a linear step over the first slice towards zero at the camera's
174    // position to recover the correct integral value.
175    let t_slice = t_max / num_slices;
176    let fade = saturate(t / t_slice);
177    // Recover the values from log space
178    return exp(sample.rgb) * fade;
179}
180
181// PHASE FUNCTIONS
182
183// -(L . V) == (L . -V). -V here is our ray direction, which points away from the view 
184// instead of towards it (which would be the *view direction*, V)
185
186// evaluates the rayleigh phase function, which describes the likelihood
187// of a rayleigh scattering event scattering light from the light direction towards the view
188fn rayleigh(neg_LdotV: f32) -> f32 {
189    return FRAC_3_16_PI * (1 + (neg_LdotV * neg_LdotV));
190}
191
192// evaluates the henyey-greenstein phase function, which describes the likelihood
193// of a mie scattering event scattering light from the light direction towards the view
194fn henyey_greenstein(neg_LdotV: f32) -> f32 {
195    let g = atmosphere.mie_asymmetry;
196    let denom = 1.0 + g * g - 2.0 * g * neg_LdotV;
197    return FRAC_4_PI * (1.0 - g * g) / (denom * sqrt(denom));
198}
199
200// ATMOSPHERE SAMPLING
201
202struct AtmosphereSample {
203    /// units: m^-1
204    rayleigh_scattering: vec3<f32>,
205
206    /// units: m^-1
207    mie_scattering: f32,
208
209    /// the sum of scattering and absorption. Since the phase function doesn't
210    /// matter for this, we combine rayleigh and mie extinction to a single 
211    //  value.
212    //
213    /// units: m^-1
214    extinction: vec3<f32>
215}
216
217/// Samples atmosphere optical densities at a given radius
218fn sample_atmosphere(r: f32) -> AtmosphereSample {
219    let altitude = clamp(r, atmosphere.bottom_radius, atmosphere.top_radius) - atmosphere.bottom_radius;
220
221    // atmosphere values at altitude
222    let mie_density = exp(-atmosphere.mie_density_exp_scale * altitude);
223    let rayleigh_density = exp(-atmosphere.rayleigh_density_exp_scale * altitude);
224    var ozone_density: f32 = max(0.0, 1.0 - (abs(altitude - atmosphere.ozone_layer_altitude) / (atmosphere.ozone_layer_width * 0.5)));
225
226    let mie_scattering = mie_density * atmosphere.mie_scattering;
227    let mie_absorption = mie_density * atmosphere.mie_absorption;
228    let mie_extinction = mie_scattering + mie_absorption;
229
230    let rayleigh_scattering = rayleigh_density * atmosphere.rayleigh_scattering;
231    // no rayleigh absorption
232    // rayleigh extinction is the sum of scattering and absorption
233
234    // ozone doesn't contribute to scattering
235    let ozone_absorption = ozone_density * atmosphere.ozone_absorption;
236
237    var sample: AtmosphereSample;
238    sample.rayleigh_scattering = rayleigh_scattering;
239    sample.mie_scattering = mie_scattering;
240    sample.extinction = rayleigh_scattering + mie_extinction + ozone_absorption;
241
242    return sample;
243}
244
245/// evaluates L_scat, equation 3 in the paper, which gives the total single-order scattering towards the view at a single point
246fn sample_local_inscattering(local_atmosphere: AtmosphereSample, ray_dir: vec3<f32>, local_r: f32, local_up: vec3<f32>) -> vec3<f32> {
247    var inscattering = vec3(0.0);
248    for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
249        let light = &lights.directional_lights[light_i];
250
251        let mu_light = dot((*light).direction_to_light, local_up);
252
253        // -(L . V) == (L . -V). -V here is our ray direction, which points away from the view
254        // instead of towards it (as is the convention for V)
255        let neg_LdotV = dot((*light).direction_to_light, ray_dir);
256
257        // Phase functions give the proportion of light
258        // scattered towards the camera for each scattering type
259        let rayleigh_phase = rayleigh(neg_LdotV);
260        let mie_phase = henyey_greenstein(neg_LdotV);
261        let scattering_coeff = local_atmosphere.rayleigh_scattering * rayleigh_phase + local_atmosphere.mie_scattering * mie_phase;
262
263        let transmittance_to_light = sample_transmittance_lut(local_r, mu_light);
264        let shadow_factor = transmittance_to_light * f32(!ray_intersects_ground(local_r, mu_light));
265
266        // Transmittance from scattering event to light source
267        let scattering_factor = shadow_factor * scattering_coeff;
268
269        // Additive factor from the multiscattering LUT
270        let psi_ms = sample_multiscattering_lut(local_r, mu_light);
271        let multiscattering_factor = psi_ms * (local_atmosphere.rayleigh_scattering + local_atmosphere.mie_scattering);
272
273        inscattering += (*light).color.rgb * (scattering_factor + multiscattering_factor);
274    }
275    return inscattering * view.exposure;
276}
277
278const SUN_ANGULAR_SIZE: f32 = 0.0174533; // angular diameter of sun in radians
279
280fn sample_sun_radiance(ray_dir_ws: vec3<f32>) -> vec3<f32> {
281    let r = view_radius();
282    let mu_view = ray_dir_ws.y;
283    let shadow_factor = f32(!ray_intersects_ground(r, mu_view));
284    var sun_radiance = vec3(0.0);
285    for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
286        let light = &lights.directional_lights[light_i];
287        let neg_LdotV = dot((*light).direction_to_light, ray_dir_ws);
288        let angle_to_sun = fast_acos(neg_LdotV);
289        let pixel_size = fwidth(angle_to_sun);
290        let factor = smoothstep(0.0, -pixel_size * ROOT_2, angle_to_sun - SUN_ANGULAR_SIZE * 0.5);
291        let sun_solid_angle = (SUN_ANGULAR_SIZE * SUN_ANGULAR_SIZE) * 4.0 * FRAC_PI;
292        sun_radiance += ((*light).color.rgb / sun_solid_angle) * factor * shadow_factor;
293    }
294    return sun_radiance;
295}
296
297// TRANSFORM UTILITIES
298
299fn max_atmosphere_distance(r: f32, mu: f32) -> f32 {
300    let t_top = distance_to_top_atmosphere_boundary(r, mu);
301    let t_bottom = distance_to_bottom_atmosphere_boundary(r, mu);
302    let hits = ray_intersects_ground(r, mu);
303    return mix(t_top, t_bottom, f32(hits));
304}
305
306/// Assuming y=0 is the planet ground, returns the view radius in meters
307fn view_radius() -> f32 {
308    return view.world_position.y * settings.scene_units_to_m + atmosphere.bottom_radius;
309}
310
311// We assume the `up` vector at the view position is the y axis, since the world is locally flat/level.
312// t = distance along view ray in atmosphere space
313// NOTE: this means that if your world is actually spherical, this will be wrong.
314fn get_local_up(r: f32, t: f32, ray_dir: vec3<f32>) -> vec3<f32> {
315    return normalize(vec3(0.0, r, 0.0) + t * ray_dir);
316}
317
318// Given a ray starting at radius r, with mu = cos(zenith angle),
319// and a t = distance along the ray, gives the new radius at point t
320fn get_local_r(r: f32, mu: f32, t: f32) -> f32 {
321    return sqrt(t * t + 2.0 * r * mu * t + r * r);
322}
323
324// Convert uv [0.0 .. 1.0] coordinate to ndc space xy [-1.0 .. 1.0]
325fn uv_to_ndc(uv: vec2<f32>) -> vec2<f32> {
326    return uv * vec2(2.0, -2.0) + vec2(-1.0, 1.0);
327}
328
329/// Convert ndc space xy coordinate [-1.0 .. 1.0] to uv [0.0 .. 1.0]
330fn ndc_to_uv(ndc: vec2<f32>) -> vec2<f32> {
331    return ndc * vec2(0.5, -0.5) + vec2(0.5);
332}
333
334/// Converts a direction in world space to atmosphere space
335fn direction_world_to_atmosphere(dir_ws: vec3<f32>) -> vec3<f32> {
336    let dir_as = atmosphere_transforms.atmosphere_from_world * vec4(dir_ws, 0.0);
337    return dir_as.xyz;
338}
339
340/// Converts a direction in atmosphere space to world space
341fn direction_atmosphere_to_world(dir_as: vec3<f32>) -> vec3<f32> {
342    let dir_ws = atmosphere_transforms.world_from_atmosphere * vec4(dir_as, 0.0);
343    return dir_ws.xyz;
344}
345
346// Modified from skybox.wgsl. For this pass we don't need to apply a separate sky transform or consider camera viewport.
347// w component is the cosine of the view direction with the view forward vector, to correct step distance at the edges of the viewport
348fn uv_to_ray_direction(uv: vec2<f32>) -> vec4<f32> {
349    // Using world positions of the fragment and camera to calculate a ray direction
350    // breaks down at large translations. This code only needs to know the ray direction.
351    // The ray direction is along the direction from the camera to the fragment position.
352    // In view space, the camera is at the origin, so the view space ray direction is
353    // along the direction of the fragment position - (0,0,0) which is just the
354    // fragment position.
355    // Use the position on the near clipping plane to avoid -inf world position
356    // because the far plane of an infinite reverse projection is at infinity.
357    let view_position_homogeneous = view.view_from_clip * vec4(
358        uv_to_ndc(uv),
359        1.0,
360        1.0,
361    );
362
363    let view_ray_direction = view_position_homogeneous.xyz / view_position_homogeneous.w;
364    // Transforming the view space ray direction by the inverse view matrix, transforms the
365    // direction to world space. Note that the w element is set to 0.0, as this is a
366    // vector direction, not a position, That causes the matrix multiplication to ignore
367    // the translations from the view matrix.
368    let ray_direction = (view.world_from_view * vec4(view_ray_direction, 0.0)).xyz;
369
370    return vec4(normalize(ray_direction), -view_ray_direction.z);
371}
372
373fn zenith_azimuth_to_ray_dir(zenith: f32, azimuth: f32) -> vec3<f32> {
374    let sin_zenith = sin(zenith);
375    let mu = cos(zenith);
376    let sin_azimuth = sin(azimuth);
377    let cos_azimuth = cos(azimuth);
378    return vec3(sin_azimuth * sin_zenith, mu, -cos_azimuth * sin_zenith);
379}