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, ray_sphere_intersect}
4
5#import bevy_pbr::atmosphere::{
6    types::Atmosphere,
7    bindings::{
8        atmosphere, settings, view, lights, transmittance_lut, atmosphere_lut_sampler,
9        multiscattering_lut, sky_view_lut, aerial_view_lut, atmosphere_transforms,
10        medium_density_lut, medium_scattering_lut, medium_sampler,
11    },
12    bruneton_functions::{
13        transmittance_lut_r_mu_to_uv, ray_intersects_ground,
14        distance_to_top_atmosphere_boundary, distance_to_bottom_atmosphere_boundary
15    },
16}
17
18// NOTE FOR CONVENTIONS: 
19// r:
20//   radius, or distance from planet center 
21//
22// altitude:
23//   distance from planet **surface**
24//
25// mu:
26//   cosine of the zenith angle of a ray with
27//   respect to the planet normal
28//
29// atmosphere space:
30//   abbreviated as "as" (contrast with vs, cs, ws), this space is similar
31//   to view space, but with the camera positioned horizontally on the planet
32//   surface, so the horizon is a horizontal line centered vertically in the
33//   frame. This enables the non-linear latitude parametrization the paper uses 
34//   to concentrate detail near the horizon 
35
36
37// CONSTANTS
38const FRAC_PI: f32 = 0.3183098862; // 1 / π
39const FRAC_2_PI: f32 = 0.15915494309;  // 1 / (2π)
40const FRAC_3_16_PI: f32 = 0.0596831036594607509; // 3 / (16π)
41const FRAC_4_PI: f32 = 0.07957747154594767; // 1 / (4π)
42const ROOT_2: f32 = 1.41421356; // √2
43const EPSILON: f32 = 1.0; // 1 meter
44const MIN_EXTINCTION: vec3<f32> = vec3(1e-12);
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 PARAMETERIZATIONS
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(atmosphere, r, mu);
118    return textureSampleLevel(transmittance_lut, atmosphere_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, atmosphere_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, atmosphere_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, atmosphere_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// ATMOSPHERE SAMPLING
182
183const ABSORPTION_DENSITY: f32 = 0.0;
184const SCATTERING_DENSITY: f32 = 1.0;
185
186// samples from the atmosphere density LUT.
187//
188// calling with `component = 0.0` will return the atmosphere's absorption density,
189// while calling with `component = 1.0` will return the atmosphere's scattering density.
190fn sample_density_lut(r: f32, component: f32) -> vec3<f32> {
191    // sampler clamps to [0, 1] anyways, no need to clamp the altitude
192    let normalized_altitude = (r - atmosphere.bottom_radius) / (atmosphere.top_radius - atmosphere.bottom_radius);
193    let uv = vec2(1.0 - normalized_altitude, component);
194    return textureSampleLevel(medium_density_lut, medium_sampler, uv, 0.0).xyz;
195}
196
197// samples from the atmosphere scattering LUT. `neg_LdotV` is the dot product
198// of the light direction and the incoming view vector.
199fn sample_scattering_lut(r: f32, neg_LdotV: f32) -> vec3<f32> {
200    let normalized_altitude = (r - atmosphere.bottom_radius) / (atmosphere.top_radius - atmosphere.bottom_radius);
201    let uv = vec2(1.0 - normalized_altitude, neg_LdotV * 0.5 + 0.5);
202    return textureSampleLevel(medium_scattering_lut, medium_sampler, uv, 0.0).xyz;
203}
204
205/// evaluates L_scat, equation 3 in the paper, which gives the total single-order scattering towards the view at a single point
206fn sample_local_inscattering(local_scattering: vec3<f32>, ray_dir: vec3<f32>, world_pos: vec3<f32>) -> vec3<f32> {
207    let local_r = length(world_pos);
208    let local_up = normalize(world_pos);
209    var inscattering = vec3(0.0);
210    for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
211        let light = &lights.directional_lights[light_i];
212
213        let mu_light = dot((*light).direction_to_light, local_up);
214
215        // -(L . V) == (L . -V). -V here is our ray direction, which points away from the view
216        // instead of towards it (as is the convention for V)
217        let neg_LdotV = dot((*light).direction_to_light, ray_dir);
218
219        let transmittance_to_light = sample_transmittance_lut(local_r, mu_light);
220        let shadow_factor = transmittance_to_light * f32(!ray_intersects_ground(local_r, mu_light));
221        let scattering_coeff = sample_scattering_lut(local_r, neg_LdotV);
222
223        // Transmittance from scattering event to light source
224        let scattering_factor = shadow_factor * scattering_coeff;
225
226        // Additive factor from the multiscattering LUT
227        let psi_ms = sample_multiscattering_lut(local_r, mu_light);
228        let multiscattering_factor = psi_ms * local_scattering;
229
230        inscattering += (*light).color.rgb * (scattering_factor + multiscattering_factor);
231    }
232    return inscattering;
233}
234
235fn sample_sun_radiance(ray_dir_ws: vec3<f32>) -> vec3<f32> {
236    let view_pos = get_view_position();
237    let r = length(view_pos);
238    let up = normalize(view_pos);
239    let mu_view = dot(ray_dir_ws, up);
240    let shadow_factor = f32(!ray_intersects_ground(r, mu_view));
241    var sun_radiance = vec3(0.0);
242    for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
243        let light = &lights.directional_lights[light_i];
244        let neg_LdotV = dot((*light).direction_to_light, ray_dir_ws);
245        let angle_to_sun = fast_acos(clamp(neg_LdotV, -1.0, 1.0));
246        let w = max(0.5 * fwidth(angle_to_sun), 1e-6);
247        let sun_angular_size = (*light).sun_disk_angular_size;
248        let sun_intensity = (*light).sun_disk_intensity;
249        if sun_angular_size > 0.0 && sun_intensity > 0.0 {
250            let factor = 1 - smoothstep(sun_angular_size * 0.5 - w, sun_angular_size * 0.5 + w, angle_to_sun);
251            let sun_solid_angle = (sun_angular_size * sun_angular_size) * 0.25 * PI;
252            sun_radiance += ((*light).color.rgb / sun_solid_angle) * sun_intensity * factor * shadow_factor;
253        }
254    }
255    return sun_radiance;
256}
257
258fn calculate_visible_sun_ratio(atmosphere: Atmosphere, r: f32, mu: f32, sun_angular_size: f32) -> f32 {
259    let bottom_radius = atmosphere.bottom_radius;
260    // Calculate the angle between horizon and sun center
261    // Invert the horizon angle calculation to fix shading direction
262    let horizon_cos = -sqrt(1.0 - (bottom_radius * bottom_radius) / (r * r));
263    let horizon_angle = fast_acos_4(horizon_cos);
264    let sun_zenith_angle = fast_acos_4(mu);
265    
266    // If sun is completely above horizon
267    if sun_zenith_angle + sun_angular_size * 0.5 <= horizon_angle {
268        return 1.0;
269    }
270    
271    // If sun is completely below horizon
272    if sun_zenith_angle - sun_angular_size * 0.5 >= horizon_angle {
273        return 0.0;
274    }
275    
276    // Calculate partial visibility using circular segment area formula
277    let d = (horizon_angle - sun_zenith_angle) / (sun_angular_size * 0.5);
278    let visible_ratio = 0.5 + d * 0.5;
279    return clamp(visible_ratio, 0.0, 1.0);
280}
281
282// TRANSFORM UTILITIES
283
284/// Clamp a position to the planet surface (with a small epsilon) to avoid underground artifacts.
285fn clamp_to_surface(atmosphere: Atmosphere, position: vec3<f32>) -> vec3<f32> {
286    let min_radius = atmosphere.bottom_radius + EPSILON;
287    let r = length(position);
288    if r < min_radius {
289        let up = normalize(position);
290        return up * min_radius;
291    }
292    return position;
293}
294
295fn max_atmosphere_distance(r: f32, mu: f32) -> f32 {
296    let t_top = distance_to_top_atmosphere_boundary(atmosphere, r, mu);
297    let t_bottom = distance_to_bottom_atmosphere_boundary(r, mu);
298    let hits = ray_intersects_ground(r, mu);
299    return mix(t_top, t_bottom, f32(hits));
300}
301
302/// Returns the observer's position in the atmosphere
303fn get_view_position() -> vec3<f32> {
304    var world_pos = view.world_position * settings.scene_units_to_m + vec3(0.0, atmosphere.bottom_radius, 0.0);
305    return clamp_to_surface(atmosphere, world_pos);
306}
307
308// We assume the `up` vector at the view position is the y axis, since the world is locally flat/level.
309// t = distance along view ray in atmosphere space
310// NOTE: this means that if your world is actually spherical, this will be wrong.
311fn get_local_up(r: f32, t: f32, ray_dir: vec3<f32>) -> vec3<f32> {
312    return normalize(vec3(0.0, r, 0.0) + t * ray_dir);
313}
314
315// Given a ray starting at radius r, with mu = cos(zenith angle),
316// and a t = distance along the ray, gives the new radius at point t
317fn get_local_r(r: f32, mu: f32, t: f32) -> f32 {
318    return sqrt(t * t + 2.0 * r * mu * t + r * r);
319}
320
321// Convert uv [0.0 .. 1.0] coordinate to ndc space xy [-1.0 .. 1.0]
322fn uv_to_ndc(uv: vec2<f32>) -> vec2<f32> {
323    return uv * vec2(2.0, -2.0) + vec2(-1.0, 1.0);
324}
325
326/// Convert ndc space xy coordinate [-1.0 .. 1.0] to uv [0.0 .. 1.0]
327fn ndc_to_uv(ndc: vec2<f32>) -> vec2<f32> {
328    return ndc * vec2(0.5, -0.5) + vec2(0.5);
329}
330
331/// Converts a direction in world space to atmosphere space
332fn direction_world_to_atmosphere(dir_ws: vec3<f32>, up: vec3<f32>) -> vec3<f32> {
333    // Camera forward in world space (-Z in view to world transform)
334    let forward_ws = (view.world_from_view * vec4(0.0, 0.0, -1.0, 0.0)).xyz;
335    let tangent_z = normalize(up * dot(forward_ws, up) - forward_ws);
336    let tangent_x = cross(up, tangent_z);
337    return vec3(
338        dot(dir_ws, tangent_x),
339        dot(dir_ws, up),
340        dot(dir_ws, tangent_z),
341    );
342}
343
344/// Converts a direction in atmosphere space to world space
345fn direction_atmosphere_to_world(dir_as: vec3<f32>) -> vec3<f32> {
346    let dir_ws = atmosphere_transforms.world_from_atmosphere * vec4(dir_as, 0.0);
347    return dir_ws.xyz;
348}
349
350// Modified from skybox.wgsl. For this pass we don't need to apply a separate sky transform or consider camera viewport.
351// Returns a normalized ray direction in world space.
352fn uv_to_ray_direction(uv: vec2<f32>) -> vec3<f32> {
353    // Using world positions of the fragment and camera to calculate a ray direction
354    // breaks down at large translations. This code only needs to know the ray direction.
355    // The ray direction is along the direction from the camera to the fragment position.
356    // In view space, the camera is at the origin, so the view space ray direction is
357    // along the direction of the fragment position - (0,0,0) which is just the
358    // fragment position.
359    // Use the position on the near clipping plane to avoid -inf world position
360    // because the far plane of an infinite reverse projection is at infinity.
361    let view_position_homogeneous = view.view_from_clip * vec4(
362        uv_to_ndc(uv),
363        1.0,
364        1.0,
365    );
366
367    let view_ray_direction = view_position_homogeneous.xyz / view_position_homogeneous.w;
368    // Transforming the view space ray direction by the inverse view matrix, transforms the
369    // direction to world space. Note that the w element is set to 0.0, as this is a
370    // vector direction, not a position, That causes the matrix multiplication to ignore
371    // the translations from the view matrix.
372    let ray_direction = (view.world_from_view * vec4(view_ray_direction, 0.0)).xyz;
373
374    return normalize(ray_direction);
375}
376
377fn zenith_azimuth_to_ray_dir(zenith: f32, azimuth: f32) -> vec3<f32> {
378    let sin_zenith = sin(zenith);
379    let mu = cos(zenith);
380    let sin_azimuth = sin(azimuth);
381    let cos_azimuth = cos(azimuth);
382    return vec3(sin_azimuth * sin_zenith, mu, -cos_azimuth * sin_zenith);
383}
384
385struct RaymarchSegment {
386    start: f32,
387    end: f32,
388}
389
390fn get_raymarch_segment(r: f32, mu: f32) -> RaymarchSegment {
391    // Get both intersection points with atmosphere
392    let atmosphere_intersections = ray_sphere_intersect(r, mu, atmosphere.top_radius);
393    let ground_intersections = ray_sphere_intersect(r, mu, atmosphere.bottom_radius);
394
395    var segment: RaymarchSegment;
396
397    if r < atmosphere.bottom_radius {
398        // Inside planet - start from bottom of atmosphere
399        segment.start = ground_intersections.y; // Use second intersection point with ground
400        segment.end = atmosphere_intersections.y;
401    } else if r < atmosphere.top_radius {
402        // Inside atmosphere
403        segment.start = 0.0;
404        segment.end = select(atmosphere_intersections.y, ground_intersections.x, ray_intersects_ground(r, mu));
405    } else {
406        // Outside atmosphere
407        if atmosphere_intersections.x < 0.0 {
408            // No intersection with atmosphere
409            return segment;
410        }
411        // Start at atmosphere entry, end at exit or ground
412        segment.start = atmosphere_intersections.x;
413        segment.end = select(atmosphere_intersections.y, ground_intersections.x, ray_intersects_ground(r, mu));
414    }
415
416    return segment;
417}
418
419struct RaymarchResult {
420    inscattering: vec3<f32>,
421    transmittance: vec3<f32>,
422}
423
424fn raymarch_atmosphere(
425    pos: vec3<f32>,
426    ray_dir: vec3<f32>,
427    t_max: f32,
428    max_samples: u32,
429    uv: vec2<f32>,
430    ground: bool
431) -> RaymarchResult {
432    let r = length(pos);
433    let up = normalize(pos);
434    let mu = dot(ray_dir, up);
435
436    // Optimization: Reduce sample count at close proximity to the scene
437    let sample_count = mix(1.0, f32(max_samples), saturate(t_max * 0.01));
438
439    let segment = get_raymarch_segment(r, mu);
440    let t_start = segment.start;
441    var t_end = segment.end;
442
443    t_end = min(t_end, t_max);
444    let t_total = t_end - t_start;
445
446    var result: RaymarchResult;
447    result.inscattering = vec3(0.0);
448    result.transmittance = vec3(1.0);
449
450    // Skip if invalid segment
451    if t_total <= 0.0 {
452        return result;
453    }
454
455    var prev_t = t_start;
456    var optical_depth = vec3(0.0);
457    for (var s = 0.0; s < sample_count; s += 1.0) {
458        // Linear distribution from atmosphere entry to exit/ground
459        let t_i = t_start + t_total * (s + MIDPOINT_RATIO) / sample_count;
460        let dt_i = (t_i - prev_t);
461        prev_t = t_i;
462
463        let sample_pos = pos + ray_dir * t_i;
464        let local_r = length(sample_pos);
465        let local_up = normalize(sample_pos);
466
467        let absorption = sample_density_lut(local_r, ABSORPTION_DENSITY);
468        let scattering = sample_density_lut(local_r, SCATTERING_DENSITY);
469        let extinction = absorption + scattering;
470
471        let sample_optical_depth = extinction * dt_i;
472        optical_depth += sample_optical_depth;
473        let sample_transmittance = exp(-sample_optical_depth);
474
475        let inscattering = sample_local_inscattering(
476            scattering,
477            ray_dir,
478            sample_pos
479        );
480
481        let s_int = (inscattering - inscattering * sample_transmittance) / max(extinction, MIN_EXTINCTION);
482        result.inscattering += result.transmittance * s_int;
483
484        result.transmittance *= sample_transmittance;
485        if all(result.transmittance < vec3(0.001)) {
486            break;
487        }
488    }
489
490    // include reflected luminance from planet ground 
491    if ground && ray_intersects_ground(r, mu) {
492        for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
493            let light = &lights.directional_lights[light_i];
494            let light_dir = (*light).direction_to_light;
495            let light_color = (*light).color.rgb;
496            let transmittance_to_ground = exp(-optical_depth);
497            // position on the sphere and get the sphere normal (up)
498            let sphere_point = pos + ray_dir * t_end;
499            let sphere_normal = normalize(sphere_point);
500            let mu_light = dot(light_dir, sphere_normal);
501            let transmittance_to_light = sample_transmittance_lut(0.0, mu_light);
502            let light_luminance = transmittance_to_light * max(mu_light, 0.0) * light_color;
503            // Normalized Lambert BRDF
504            let ground_luminance = transmittance_to_ground * atmosphere.ground_albedo / PI;
505            result.inscattering += ground_luminance * light_luminance;
506        }
507    }
508
509    return result;
510}