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            uv_to_ndc, max_atmosphere_distance, uv_to_ray_direction, 
10            MIDPOINT_RATIO, get_view_position
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 world_pos = get_view_position();
26    
27    let r = length(world_pos);
28    let t_max = settings.aerial_view_lut_max_distance;
29
30    var prev_t = 0.0;
31    var total_inscattering = vec3(0.0);
32    var throughput = vec3(1.0);
33
34    for (var slice_i: u32 = 0; slice_i < settings.aerial_view_lut_size.z; slice_i++) {
35        for (var step_i: u32 = 0; step_i < settings.aerial_view_lut_samples; step_i++) {
36            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);
37            let dt = (t_i - prev_t);
38            prev_t = t_i;
39
40            let sample_pos = world_pos + ray_dir * t_i;
41            let local_r = length(sample_pos);
42            let local_up = normalize(sample_pos);
43
44            let local_atmosphere = sample_atmosphere(local_r);
45            let sample_optical_depth = local_atmosphere.extinction * dt;
46            let sample_transmittance = exp(-sample_optical_depth);
47
48            // evaluate one segment of the integral
49            var inscattering = sample_local_inscattering(local_atmosphere, ray_dir, sample_pos);
50
51            // Analytical integration of the single scattering term in the radiance transfer equation
52            let s_int = (inscattering - inscattering * sample_transmittance) / local_atmosphere.extinction;
53            total_inscattering += throughput * s_int;
54
55            throughput *= sample_transmittance;
56            if all(throughput < vec3(0.001)) {
57                break;
58            }
59        }
60
61        // Store in log space to allow linear interpolation of exponential values between slices
62        let log_inscattering = log(max(total_inscattering, vec3(1e-6)));
63        textureStore(aerial_view_lut_out, vec3(vec2<u32>(idx.xy), slice_i), vec4(log_inscattering, 0.0));
64    }
65}