1//! NOTE FOR CONVENTIONS:
2//! r:
3//!   radius, or distance from planet center
4//!
5//! altitude:
6//!   distance from planet **surface**
7//!
8//! mu:
9//!   cosine of the zenith angle of a ray with
10//!   respect to the planet normal
11//!
12//! atmosphere space:
13//!   abbreviated as "as" (contrast with vs, cs, ws), this space is similar
14//!   to view space, but with the camera positioned horizontally on the planet
15//!   surface, so the horizon is a horizontal line centered vertically in the
16//!   frame. This enables the non-linear latitude parametrization the paper uses
17//!   to concentrate detail near the horizon
18
19import bevy_render::maths::{PI, HALF_PI, PI_2, fast_acos, fast_acos_4, fast_atan2, ray_sphere_intersect};
20
21import package::atmosphere::{
22    types::Atmosphere,
23    bindings::{
24        atmosphere, settings, view, lights, transmittance_lut, atmosphere_lut_sampler,
25        multiscattering_lut, sky_view_lut, aerial_view_lut, atmosphere_transforms,
26        medium_density_lut, medium_scattering_lut, medium_sampler,
27    },
28    bruneton_functions::{
29        transmittance_lut_r_mu_to_uv, ray_intersects_ground,
30        distance_to_top_atmosphere_boundary, distance_to_bottom_atmosphere_boundary
31    },
32};
33
34// CONSTANTS
35/// 1 / π
36const FRAC_PI: f32 = 0.3183098862;
37/// 1 / (2π)
38const FRAC_2_PI: f32 = 0.15915494309;
39/// 3 / (16π)
40const FRAC_3_16_PI: f32 = 0.0596831036594607509;
41/// 1 / (4π)
42const FRAC_4_PI: f32 = 0.07957747154594767;
43/// √2
44const ROOT_2: f32 = 1.41421356;
45/// 1 meter
46const EPSILON: f32 = 1.0;
47const MIN_EXTINCTION: vec3<f32> = vec3(1e-12);
48
49/// During raymarching, each segment is sampled at a single point. This constant determines
50/// where in the segment that sample is taken (0.0 = start, 0.5 = middle, 1.0 = end).
51/// We use 0.3 to sample closer to the start of each segment, which better approximates
52/// the exponential falloff of atmospheric density.
53const MIDPOINT_RATIO: f32 = 0.3;
54
55// LUT UV PARAMETERIZATIONS
56
57fn unit_to_sub_uvs(val: vec2<f32>, resolution: vec2<f32>) -> vec2<f32> {
58    return (val + 0.5f / resolution) * (resolution / (resolution + 1.0f));
59}
60
61fn sub_uvs_to_unit(val: vec2<f32>, resolution: vec2<f32>) -> vec2<f32> {
62    return (val - 0.5f / resolution) * (resolution / (resolution - 1.0f));
63}
64
65fn multiscattering_lut_r_mu_to_uv(r: f32, mu: f32) -> vec2<f32> {
66    let u = 0.5 + 0.5 * mu;
67    let v = saturate((r - atmosphere.inner_radius) / (atmosphere.outer_radius - atmosphere.inner_radius)); //TODO
68    return unit_to_sub_uvs(vec2(u, v), vec2<f32>(settings.multiscattering_lut_size));
69}
70
71fn multiscattering_lut_uv_to_r_mu(uv: vec2<f32>) -> vec2<f32> {
72    let adj_uv = sub_uvs_to_unit(uv, vec2<f32>(settings.multiscattering_lut_size));
73    let r = mix(atmosphere.inner_radius, atmosphere.outer_radius, adj_uv.y);
74    let mu = adj_uv.x * 2 - 1;
75    return vec2(r, mu);
76}
77
78fn sky_view_lut_r_mu_azimuth_to_uv(r: f32, mu: f32, azimuth: f32) -> vec2<f32> {
79    let u = (azimuth * FRAC_2_PI) + 0.5;
80
81    let v_horizon = sqrt(r * r - atmosphere.inner_radius * atmosphere.inner_radius);
82    let cos_beta = v_horizon / r;
83    // Using fast_acos_4 for better precision at small angles
84    // to avoid artifacts at the horizon
85    let beta = fast_acos_4(cos_beta);
86    let horizon_zenith = PI - beta;
87    let view_zenith = fast_acos_4(mu);
88
89    // Apply non-linear transformation to compress more texels
90    // near the horizon where high-frequency details matter most
91    // l is latitude in [-π/2, π/2] and v is texture coordinate in [0,1]
92    let l = view_zenith - horizon_zenith;
93    let abs_l = abs(l);
94
95    let v = 0.5 + 0.5 * sign(l) * sqrt(abs_l / HALF_PI);
96
97    return unit_to_sub_uvs(vec2(u, v), vec2<f32>(settings.sky_view_lut_size));
98}
99
100fn sky_view_lut_uv_to_zenith_azimuth(r: f32, uv: vec2<f32>) -> vec2<f32> {
101    let adj_uv = sub_uvs_to_unit(vec2(uv.x, 1.0 - uv.y), vec2<f32>(settings.sky_view_lut_size));
102    let azimuth = (adj_uv.x - 0.5) * PI_2;
103
104    // Horizon parameters
105    let v_horizon = sqrt(r * r - atmosphere.inner_radius * atmosphere.inner_radius);
106    let cos_beta = v_horizon / r;
107    let beta = fast_acos_4(cos_beta);
108    let horizon_zenith = PI - beta;
109
110    // Inverse of horizon-detail mapping to recover original latitude from texture coordinate
111    let t = abs(2.0 * (adj_uv.y - 0.5));
112    let l = sign(adj_uv.y - 0.5) * HALF_PI * t * t;
113
114    return vec2(horizon_zenith - l, azimuth);
115}
116
117// LUT SAMPLING
118
119fn sample_transmittance_lut(r: f32, mu: f32) -> vec3<f32> {
120    let uv = transmittance_lut_r_mu_to_uv(atmosphere, r, mu);
121    return textureSampleLevel(transmittance_lut, atmosphere_lut_sampler, uv, 0.0).rgb;
122}
123
124fn sample_multiscattering_lut(r: f32, mu: f32) -> vec3<f32> {
125    let uv = multiscattering_lut_r_mu_to_uv(r, mu);
126    return textureSampleLevel(multiscattering_lut, atmosphere_lut_sampler, uv, 0.0).rgb;
127}
128
129fn sample_sky_view_lut(r: f32, ray_dir_as: vec3<f32>) -> vec3<f32> {
130    let mu = ray_dir_as.y;
131    let azimuth = fast_atan2(ray_dir_as.x, -ray_dir_as.z);
132    let uv = sky_view_lut_r_mu_azimuth_to_uv(r, mu, azimuth);
133    return textureSampleLevel(sky_view_lut, atmosphere_lut_sampler, uv, 0.0).rgb;
134}
135
136fn ndc_to_camera_dist(ndc: vec3<f32>) -> f32 {
137    let view_pos = view.view_from_clip * vec4(ndc, 1.0);
138    let p_view = view_pos.xyz / view_pos.w;
139    let p_world = (view.world_from_view * vec4(p_view, 1.0)).xyz;
140    let p_atmo = (atmosphere.world_to_atmosphere * vec4(p_world, 1.0)).xyz;
141    let cam_atmo = (atmosphere.world_to_atmosphere * vec4(view.world_position, 1.0)).xyz;
142    return length(p_atmo - cam_atmo);
143}
144
145/// RGB channels: total inscattered light along the camera ray to the current sample.
146/// A channel: average transmittance across all wavelengths to the current sample.
147fn sample_aerial_view_lut(uv: vec2<f32>, t: f32) -> vec3<f32> {
148    let t_max = settings.aerial_view_lut_max_distance;
149    let num_slices = f32(settings.aerial_view_lut_size.z);
150    // Each texel stores the value of the scattering integral over the whole slice,
151    // which requires us to offset the w coordinate by half a slice. For
152    // example, if we wanted the value of the integral at the boundary between slices,
153    // we'd need to sample at the center of the previous slice, and vice-versa for
154    // sampling in the center of a slice.
155    let uvw = vec3(uv, saturate(t / t_max - 0.5 / num_slices));
156    let sample = textureSampleLevel(aerial_view_lut, atmosphere_lut_sampler, uvw, 0.0);
157    // Since sampling anywhere between w=0 and w=t_slice will clamp to the first slice,
158    // we need to do a linear step over the first slice towards zero at the camera's
159    // position to recover the correct integral value.
160    let t_slice = t_max / num_slices;
161    let fade = saturate(t / t_slice);
162    // Recover the values from log space
163    return exp(sample.rgb) * fade;
164}
165
166// ATMOSPHERE SAMPLING
167
168const ABSORPTION_DENSITY: f32 = 0.0;
169const SCATTERING_DENSITY: f32 = 1.0;
170
171/// samples from the atmosphere density LUT.
172///
173/// calling with `component = 0.0` will return the atmosphere's absorption density,
174/// while calling with `component = 1.0` will return the atmosphere's scattering density.
175fn sample_density_lut(r: f32, component: f32) -> vec3<f32> {
176    // sampler clamps to [0, 1] anyways, no need to clamp the altitude
177    let normalized_altitude = (r - atmosphere.inner_radius) / (atmosphere.outer_radius - atmosphere.inner_radius);
178    let uv = vec2(1.0 - normalized_altitude, component);
179    return textureSampleLevel(medium_density_lut, medium_sampler, uv, 0.0).xyz;
180}
181
182
183const PHASE_MAPPING_N: f32 = 0.5;
184
185/// samples from the atmosphere scattering LUT. `neg_LdotV` is the dot product
186/// of the light direction and the incoming view vector.
187/// Nonlinear phase mapping to mitigate banding in low-resolution LUTs.
188fn sample_scattering_lut(r: f32, neg_LdotV: f32) -> vec3<f32> {
189    let normalized_altitude = (r - atmosphere.inner_radius) / (atmosphere.outer_radius - atmosphere.inner_radius);
190    let phase_uv = 0.5 + 0.5 * sign(neg_LdotV) * (1.0 - pow(1.0 - abs(neg_LdotV), PHASE_MAPPING_N));
191    let uv = vec2(1.0 - normalized_altitude, phase_uv);
192    return textureSampleLevel(medium_scattering_lut, medium_sampler, uv, 0.0).xyz;
193}
194
195/// evaluates L_scat, equation 3 in the paper, which gives the total single-order scattering towards the view at a single point
196fn sample_local_inscattering(local_scattering: vec3<f32>, ray_dir: vec3<f32>, world_pos: vec3<f32>) -> vec3<f32> {
197    let local_r = length(world_pos);
198    let local_up = normalize(world_pos);
199    var inscattering = vec3(0.0);
200    for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
201        let light = &lights.directional_lights[light_i];
202
203        let mu_light = dot((*light).direction_to_light, local_up);
204
205        // -(L . V) == (L . -V). -V here is our ray direction, which points away from the view
206        // instead of towards it (as is the convention for V)
207        let neg_LdotV = dot((*light).direction_to_light, ray_dir);
208
209        let transmittance_to_light = sample_transmittance_lut(local_r, mu_light);
210        let shadow_factor = transmittance_to_light * f32(!ray_intersects_ground(local_r, mu_light));
211        let scattering_coeff = sample_scattering_lut(local_r, neg_LdotV);
212
213        // Transmittance from scattering event to light source
214        let scattering_factor = shadow_factor * scattering_coeff;
215
216        // Additive factor from the multiscattering LUT
217        let psi_ms = sample_multiscattering_lut(local_r, mu_light);
218        let multiscattering_factor = psi_ms * local_scattering;
219
220        inscattering += (*light).color.rgb * (scattering_factor + multiscattering_factor);
221    }
222    return inscattering;
223}
224
225fn sample_sun_radiance(ray_dir_ws: vec3<f32>) -> vec3<f32> {
226    let view_pos = get_view_position();
227    let r = length(view_pos);
228    let up = normalize(view_pos);
229    let mu_view = dot(ray_dir_ws, up);
230    let shadow_factor = f32(!ray_intersects_ground(r, mu_view));
231    var sun_radiance = vec3(0.0);
232    for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
233        let light = &lights.directional_lights[light_i];
234        let neg_LdotV = dot((*light).direction_to_light, ray_dir_ws);
235        let angle_to_sun = fast_acos(clamp(neg_LdotV, -1.0, 1.0));
236        let w = max(0.5 * fwidth(angle_to_sun), 1e-6);
237        let sun_angular_size = (*light).sun_disk_angular_size;
238        let sun_intensity = (*light).sun_disk_intensity;
239        if sun_angular_size > 0.0 && sun_intensity > 0.0 {
240            let factor = 1 - smoothstep(sun_angular_size * 0.5 - w, sun_angular_size * 0.5 + w, angle_to_sun);
241            let sun_solid_angle = (sun_angular_size * sun_angular_size) * 0.25 * PI;
242            sun_radiance += ((*light).color.rgb / sun_solid_angle) * sun_intensity * factor * shadow_factor;
243        }
244    }
245    return sun_radiance;
246}
247
248fn calculate_visible_sun_ratio(atmosphere: Atmosphere, r: f32, mu: f32, sun_angular_size: f32) -> f32 {
249    let inner_radius = atmosphere.inner_radius;
250    // Calculate the angle between horizon and sun center
251    // Invert the horizon angle calculation to fix shading direction
252    let horizon_cos = -sqrt(1.0 - (inner_radius * inner_radius) / (r * r));
253    let horizon_angle = fast_acos_4(horizon_cos);
254    let sun_zenith_angle = fast_acos_4(mu);
255
256    // If sun is completely above horizon
257    if sun_zenith_angle + sun_angular_size * 0.5 <= horizon_angle {
258        return 1.0;
259    }
260
261    // If sun is completely below horizon
262    if sun_zenith_angle - sun_angular_size * 0.5 >= horizon_angle {
263        return 0.0;
264    }
265
266    // Calculate partial visibility using circular segment area formula
267    let d = (horizon_angle - sun_zenith_angle) / (sun_angular_size * 0.5);
268    let visible_ratio = 0.5 + d * 0.5;
269    return clamp(visible_ratio, 0.0, 1.0);
270}
271
272// TRANSFORM UTILITIES
273
274/// Clamp a position to the planet surface (with a small epsilon) to avoid underground artifacts.
275fn clamp_to_surface(atmosphere: Atmosphere, position: vec3<f32>) -> vec3<f32> {
276    let min_radius = atmosphere.inner_radius + EPSILON;
277    let r = length(position);
278    if r < min_radius {
279        let up = normalize(position);
280        return up * min_radius;
281    }
282    return position;
283}
284
285fn max_atmosphere_distance(r: f32, mu: f32) -> f32 {
286    let t_top = distance_to_top_atmosphere_boundary(atmosphere, r, mu);
287    let t_bottom = distance_to_bottom_atmosphere_boundary(r, mu);
288    let hits = ray_intersects_ground(r, mu);
289    return mix(t_top, t_bottom, f32(hits));
290}
291
292/// Returns the observer's position in the atmosphere
293fn get_view_position() -> vec3<f32> {
294    let atmo_pos = (atmosphere.world_to_atmosphere * vec4(view.world_position, 1.0)).xyz;
295    return clamp_to_surface(atmosphere, atmo_pos);
296}
297
298/// We assume the `up` vector at the view position is the y axis, since the world is locally flat/level.
299/// t = distance along view ray in atmosphere space
300/// NOTE: this means that if your world is actually spherical, this will be wrong.
301fn get_local_up(r: f32, t: f32, ray_dir: vec3<f32>) -> vec3<f32> {
302    return normalize(vec3(0.0, r, 0.0) + t * ray_dir);
303}
304
305/// Given a ray starting at radius r, with mu = cos(zenith angle),
306/// and a t = distance along the ray, gives the new radius at point t
307fn get_local_r(r: f32, mu: f32, t: f32) -> f32 {
308    return sqrt(t * t + 2.0 * r * mu * t + r * r);
309}
310
311/// Convert uv [0.0 .. 1.0] coordinate to ndc space xy [-1.0 .. 1.0]
312fn uv_to_ndc(uv: vec2<f32>) -> vec2<f32> {
313    return uv * vec2(2.0, -2.0) + vec2(-1.0, 1.0);
314}
315
316/// Convert ndc space xy coordinate [-1.0 .. 1.0] to uv [0.0 .. 1.0]
317fn ndc_to_uv(ndc: vec2<f32>) -> vec2<f32> {
318    return ndc * vec2(0.5, -0.5) + vec2(0.5);
319}
320
321/// Converts a direction in world space to atmosphere space
322fn direction_world_to_atmosphere(dir_ws: vec3<f32>) -> vec3<f32> {
323    let dir_as = atmosphere_transforms.atmosphere_from_world * vec4(dir_ws, 0.0);
324    return dir_as.xyz;
325}
326
327/// Converts a direction in atmosphere space to world space
328fn direction_atmosphere_to_world(dir_as: vec3<f32>) -> vec3<f32> {
329    let dir_ws = atmosphere_transforms.world_from_atmosphere * vec4(dir_as, 0.0);
330    return dir_ws.xyz;
331}
332
333// Modified from skybox.wesl. For this pass we don't need to apply a separate sky transform or consider camera viewport.
334/// Returns a normalized ray direction in world space.
335fn uv_to_ray_direction(uv: vec2<f32>) -> vec3<f32> {
336    // Using world positions of the fragment and camera to calculate a ray direction
337    // breaks down at large translations. This code only needs to know the ray direction.
338    // The ray direction is along the direction from the camera to the fragment position.
339    // In view space, the camera is at the origin, so the view space ray direction is
340    // along the direction of the fragment position - (0,0,0) which is just the
341    // fragment position.
342    // Use the position on the near clipping plane to avoid -inf world position
343    // because the far plane of an infinite reverse projection is at infinity.
344    let view_position_homogeneous = view.view_from_clip * vec4(
345        uv_to_ndc(uv),
346        1.0,
347        1.0,
348    );
349
350    let view_ray_direction = view_position_homogeneous.xyz / view_position_homogeneous.w;
351    // Transforming the view space ray direction by the inverse view matrix, transforms the
352    // direction to world space. Note that the w element is set to 0.0, as this is a
353    // vector direction, not a position, That causes the matrix multiplication to ignore
354    // the translations from the view matrix.
355    let ray_direction = (view.world_from_view * vec4(view_ray_direction, 0.0)).xyz;
356
357    return normalize(ray_direction);
358}
359
360fn zenith_azimuth_to_ray_dir(zenith: f32, azimuth: f32) -> vec3<f32> {
361    let sin_zenith = sin(zenith);
362    let mu = cos(zenith);
363    let sin_azimuth = sin(azimuth);
364    let cos_azimuth = cos(azimuth);
365    return vec3(sin_azimuth * sin_zenith, mu, -cos_azimuth * sin_zenith);
366}
367
368struct RaymarchSegment {
369    start: f32,
370    end: f32,
371}
372
373fn get_raymarch_segment(r: f32, mu: f32) -> RaymarchSegment {
374    // Get both intersection points with atmosphere
375    let atmosphere_intersections = ray_sphere_intersect(r, mu, atmosphere.outer_radius);
376    let ground_intersections = ray_sphere_intersect(r, mu, atmosphere.inner_radius);
377
378    var segment: RaymarchSegment;
379
380    if r < atmosphere.inner_radius {
381        // Inside planet - start from bottom of atmosphere
382        segment.start = ground_intersections.y; // Use second intersection point with ground
383        segment.end = atmosphere_intersections.y;
384    } else if r < atmosphere.outer_radius {
385        // Inside atmosphere
386        segment.start = 0.0;
387        segment.end = select(atmosphere_intersections.y, ground_intersections.x, ray_intersects_ground(r, mu));
388    } else {
389        // Outside atmosphere
390        if atmosphere_intersections.x < 0.0 {
391            // No intersection with atmosphere
392            return segment;
393        }
394        // Start at atmosphere entry, end at exit or ground
395        segment.start = atmosphere_intersections.x;
396        segment.end = select(atmosphere_intersections.y, ground_intersections.x, ray_intersects_ground(r, mu));
397    }
398
399    return segment;
400}
401
402struct RaymarchResult {
403    inscattering: vec3<f32>,
404    transmittance: vec3<f32>,
405}
406
407fn raymarch_atmosphere(
408    pos: vec3<f32>,
409    ray_dir: vec3<f32>,
410    t_max: f32,
411    max_samples: u32,
412    uv: vec2<f32>,
413    ground: bool
414) -> RaymarchResult {
415    let r = length(pos);
416    let up = normalize(pos);
417    let mu = dot(ray_dir, up);
418
419    // Optimization: Reduce sample count at close proximity to the scene
420    let sample_count = mix(1.0, f32(max_samples), saturate(t_max * 0.01));
421
422    let segment = get_raymarch_segment(r, mu);
423    let t_start = segment.start;
424    var t_end = segment.end;
425
426    t_end = min(t_end, t_max);
427    let t_total = t_end - t_start;
428
429    var result: RaymarchResult;
430    result.inscattering = vec3(0.0);
431    result.transmittance = vec3(1.0);
432
433    // Skip if invalid segment
434    if t_total <= 0.0 {
435        return result;
436    }
437
438    var prev_t = t_start;
439    var optical_depth = vec3(0.0);
440    for (var s = 0.0; s < sample_count; s += 1.0) {
441        // Linear distribution from atmosphere entry to exit/ground
442        let t_i = t_start + t_total * (s + MIDPOINT_RATIO) / sample_count;
443        let dt_i = (t_i - prev_t);
444        prev_t = t_i;
445
446        let sample_pos = pos + ray_dir * t_i;
447        let local_r = length(sample_pos);
448        let local_up = normalize(sample_pos);
449
450        let absorption = sample_density_lut(local_r, ABSORPTION_DENSITY);
451        let scattering = sample_density_lut(local_r, SCATTERING_DENSITY);
452        let extinction = absorption + scattering;
453
454        let sample_optical_depth = extinction * dt_i;
455        optical_depth += sample_optical_depth;
456        let sample_transmittance = exp(-sample_optical_depth);
457
458        let inscattering = sample_local_inscattering(
459            scattering,
460            ray_dir,
461            sample_pos
462        );
463
464        let s_int = (inscattering - inscattering * sample_transmittance) / max(extinction, MIN_EXTINCTION);
465        result.inscattering += result.transmittance * s_int;
466
467        result.transmittance *= sample_transmittance;
468        if all(result.transmittance < vec3(0.001)) {
469            break;
470        }
471    }
472
473    // include reflected luminance from planet ground
474    if ground && ray_intersects_ground(r, mu) {
475        for (var light_i: u32 = 0u; light_i < lights.n_directional_lights; light_i++) {
476            let light = &lights.directional_lights[light_i];
477            let light_dir = (*light).direction_to_light;
478            let light_color = (*light).color.rgb;
479            let transmittance_to_ground = exp(-optical_depth);
480            // position on the sphere and get the sphere normal (up)
481            let sphere_point = pos + ray_dir * t_end;
482            let sphere_normal = normalize(sphere_point);
483            let mu_light = dot(light_dir, sphere_normal);
484            let transmittance_to_light = sample_transmittance_lut(0.0, mu_light);
485            let light_luminance = transmittance_to_light * max(mu_light, 0.0) * light_color;
486            // Normalized Lambert BRDF
487            let ground_luminance = transmittance_to_ground * atmosphere.ground_albedo / PI;
488            result.inscattering += ground_luminance * light_luminance;
489        }
490    }
491
492    return result;
493}