1#import bevy_pbr::atmosphere::{
2    bindings::settings,
3    functions::{
4        sample_density_lut, get_local_r, max_atmosphere_distance,
5        MIDPOINT_RATIO, ABSORPTION_DENSITY, SCATTERING_DENSITY
6    },
7    bruneton_functions::transmittance_lut_uv_to_r_mu,
8}
9
10#import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput
11
12@group(0) @binding(13) var transmittance_lut_out: texture_storage_2d<rgba16float, write>;
13
14@compute 
15@workgroup_size(16, 16, 1)
16fn main(@builtin(global_invocation_id) idx: vec3<u32>) {
17    let uv: vec2<f32> = (vec2<f32>(idx.xy) + 0.5) / vec2<f32>(settings.transmittance_lut_size);
18    // map UV coordinates to view height (r) and zenith cos angle (mu)
19    let r_mu = transmittance_lut_uv_to_r_mu(uv);
20
21    // compute the optical depth from view height r to the top atmosphere boundary
22    let optical_depth = ray_optical_depth(r_mu.x, r_mu.y, settings.transmittance_lut_samples);
23    let transmittance = exp(-optical_depth);
24
25    textureStore(transmittance_lut_out, idx.xy, vec4(transmittance, 1.0));
26}
27
28/// Compute the optical depth of the atmosphere from the ground to the top atmosphere boundary
29/// at a given view height (r) and zenith cos angle (mu)
30fn ray_optical_depth(r: f32, mu: f32, sample_count: u32) -> vec3<f32> {
31    let t_max = max_atmosphere_distance(r, mu);
32    var optical_depth = vec3<f32>(0.0f);
33    var prev_t = 0.0f;
34
35    for (var i = 0u; i < sample_count; i++) {
36        let t_i = t_max * (f32(i) + MIDPOINT_RATIO) / f32(sample_count);
37        let dt = t_i - prev_t;
38        prev_t = t_i;
39
40        let r_i = get_local_r(r, mu, t_i);
41
42        // PERF: A possible later optimization would be to sample at `component = 0.5`
43        // (getting the average of the two rows) and then multiplying by 2 to find the sum. 
44        let absorption = sample_density_lut(r_i, ABSORPTION_DENSITY);
45        let scattering = sample_density_lut(r_i, SCATTERING_DENSITY);
46        let extinction = absorption + scattering;
47        optical_depth += extinction * dt;
48    }
49
50    return optical_depth;
51}