1#import bevy_pbr::{
2    mesh_view_types::{Lights, DirectionalLight},
3    atmosphere::{
4        types::{Atmosphere, AtmosphereSettings},
5        bindings::{atmosphere, settings, view, lights, aerial_view_lut_out},
6        functions::{
7            sample_transmittance_lut, sample_atmosphere, rayleigh, henyey_greenstein,
8            sample_multiscattering_lut, AtmosphereSample, sample_local_inscattering,
9            get_local_r, get_local_up, view_radius, uv_to_ndc, max_atmosphere_distance,
10            uv_to_ray_direction, MIDPOINT_RATIO
11        },
12    }
13}
14
15
16@group(0) @binding(13) var aerial_view_lut_out: texture_storage_3d<rgba16float, write>;
17
18@compute
19@workgroup_size(16, 16, 1)
20fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
21    if any(idx.xy > settings.aerial_view_lut_size.xy) { return; }
22
23    let uv = (vec2<f32>(idx.xy) + 0.5) / vec2<f32>(settings.aerial_view_lut_size.xy);
24    let ray_dir = uv_to_ray_direction(uv);
25    let r = view_radius();
26    let mu = ray_dir.y;
27    let t_max = settings.aerial_view_lut_max_distance;
28
29    var prev_t = 0.0;
30    var total_inscattering = vec3(0.0);
31    var throughput = vec3(1.0);
32
33    for (var slice_i: u32 = 0; slice_i < settings.aerial_view_lut_size.z; slice_i++) {
34        for (var step_i: u32 = 0; step_i < settings.aerial_view_lut_samples; step_i++) {
35            let t_i = t_max * (f32(slice_i) + ((f32(step_i) + MIDPOINT_RATIO) / f32(settings.aerial_view_lut_samples))) / f32(settings.aerial_view_lut_size.z);
36            let dt = (t_i - prev_t);
37            prev_t = t_i;
38
39            let local_r = get_local_r(r, mu, t_i);
40            let local_up = get_local_up(r, t_i, ray_dir.xyz);
41
42            let local_atmosphere = sample_atmosphere(local_r);
43            let sample_optical_depth = local_atmosphere.extinction * dt;
44            let sample_transmittance = exp(-sample_optical_depth);
45
46            // evaluate one segment of the integral
47            var inscattering = sample_local_inscattering(local_atmosphere, ray_dir.xyz, local_r, local_up);
48
49            // Analytical integration of the single scattering term in the radiance transfer equation
50            let s_int = (inscattering - inscattering * sample_transmittance) / local_atmosphere.extinction;
51            total_inscattering += throughput * s_int;
52
53            throughput *= sample_transmittance;
54            if all(throughput < vec3(0.001)) {
55                break;
56            }
57        }
58
59        // Store in log space to allow linear interpolation of exponential values between slices
60        let log_inscattering = log(max(total_inscattering, vec3(1e-6)));
61        textureStore(aerial_view_lut_out, vec3(vec2<u32>(idx.xy), slice_i), vec4(log_inscattering, 0.0));
62    }
63}